Agent System

from langgraph.graph import StateGraph, START, END

workflow = StateGraph(State)
workflow.add_node("agentic_tool_flow", agent.agentic_tool_flow)
workflow.add_edge(START, "agentic_tool_flow")
workflow.add_edge("agentic_tool_flow", END)
app = workflow.compile(checkpointer=memory)

class State(TypedDict):
    messages: List[BaseMessage]      # Chat history
    user_id: str                     # Authenticated user
    session_id: str                  # Chat session ID
    question: str                    # User query
    provider_preference: List[str]   # Connected cloud providers
    selected_project_id: str         # Active project/subscription
    mode: str                        # "agent" or "ask" mode
    attachments: List[Dict]          # File uploads
    session_files: List[Dict]        # S3 storage files

from chat.backend.agent.utils.safe_memory_saver import SafeMemorySaver

memory = SafeMemorySaver()
config = RunnableConfig({"configurable": {"thread_id": session_id}})
app = workflow.compile(checkpointer=memory)

{
  "type": "init",
  "user_id": "user-123"
}

{
  "type": "query",
  "query": "List all GKE clusters in my project",
  "user_id": "user-123",
  "session_id": "session-456",
  "mode": "agent",
  "model": "openai/gpt-4",
  "provider_preference": ["gcp"]
}

from chat.backend.agent.tools.cloud_tools import set_user_context

set_user_context(
    user_id=state.user_id,
    session_id=state.session_id,
    provider_preference=provider_preference,
    selected_project_id=selected_project_id,
    state=state,
    mode=mode
)

from chat.backend.agent.prompt.prompt_builder import build_prompt_segments

segments = build_prompt_segments(
    provider_preference=provider_preference,
    mode=mode,
    has_zip_reference=has_zip,
    state=state
)

system_prompt = assemble_system_prompt(segments)

from chat.backend.agent.tools.cloud_tools import get_cloud_tools

tools = get_cloud_tools(
    provider_preference=provider_preference,
    mode=mode,
    state=state
)

async for event_type, event_data in workflow.stream(state):
    if event_type == "token":
        # Stream LLM tokens to frontend
        await websocket.send(json.dumps({
            "type": "message",
            "data": {"text": token_text, "is_chunk": True}
        }))
    elif event_type == "messages":
        # Handle tool calls from LLM
        if msg_chunk.tool_calls:
            for tool_call in msg_chunk.tool_calls:
                # Execute tool and return result

@tool
def list_gke_clusters(project_id: str, location: str = "-") -> str:
    """List all GKE clusters in a GCP project."""
    # Tool execution logic
    return json.dumps({"status": "success", "clusters": [...]})

{
  "type": "message",
  "data": {
    "text": "Found 3 GKE clusters:\n1. prod-cluster-1\n2. staging-cluster\n3. dev-cluster",
    "is_chunk": false,
    "is_complete": true
  }
}

from langchain.agents import tool
from chat.backend.agent.tools.cloud_tools import get_user_context

@tool
def list_gke_clusters(project_id: str, location: str = "-") -> str:
    """List all GKE clusters in a GCP project.
    
    Args:
        project_id: GCP project ID
        location: GCP region or zone (default: all)
    
    Returns:
        JSON string with cluster details
    """
    # Get user context
    ctx = get_user_context()
    user_id = ctx["user_id"]
    
    # Execute operation
    from connectors.gcp_connector.gcp.projects import list_gke_clusters
    clusters = list_gke_clusters(user_id, project_id, location)
    
    return json.dumps({"status": "success", "clusters": clusters})

from chat.backend.agent.tools.cloud_tools import get_user_context

ctx = get_user_context()
user_id = ctx["user_id"]
session_id = ctx["session_id"]
provider = ctx["provider_preference"]
mode = ctx["mode"]  # "agent" or "ask"

from chat.backend.agent.tools.cloud_tools import get_websocket_sender

sender = get_websocket_sender()
if sender:
    await sender({
        "type": "tool_progress",
        "data": {"message": "Deploying cluster..."}
    })

# Success
return json.dumps({
    "status": "success",
    "data": {...}
})

# Failure
return json.dumps({
    "status": "failed",
    "error": "Permission denied",
    "message": "User lacks IAM permissions"
})

from chat.backend.agent.access import ModeAccessController

if not ModeAccessController.is_tool_allowed(mode, tool_name):
    return {"error": "Tool not available in Ask mode"}

from chat.backend.agent.llm import LLMManager

llm_manager = LLMManager()
llm = llm_manager.get_llm(
    model="openai/gpt-4",
    streaming=True,
    temperature=0.0
)

from chat.backend.agent.utils.prefix_cache import PrefixCacheManager

cache_manager = PrefixCacheManager()
cache_manager.register_breakpoint(
    "system_prompt",
    content=system_prompt,
    priority=100
)

from chat.backend.agent.tools.mcp_preloader import start_mcp_preloader

mcp_preloader = start_mcp_preloader()

from chat.backend.agent.tools.mcp_tools import get_mcp_tools_for_user

mcp_tools = get_mcp_tools_for_user(user_id)
all_tools = base_tools + mcp_tools

async for event_type, event_data in workflow.stream(state):
    if event_type == "token":
        await websocket.send(json.dumps({
            "type": "message",
            "data": {"text": token, "is_chunk": True}
        }))

# Tool call started
await websocket.send(json.dumps({
    "type": "tool_call",
    "data": {
        "tool_name": "list_gke_clusters",
        "input": {"project_id": "my-project"},
        "status": "running"
    }
}))

# Tool call completed
await websocket.send(json.dumps({
    "type": "tool_result",
    "data": {
        "tool_name": "list_gke_clusters",
        "result": {...},
        "status": "success"
    }
}))

{
  "type": "control",
  "action": "cancel",
  "session_id": "session-456"
}

from utils.terraform.terraform_cleanup import cleanup_terraform_directory

cleanup_terraform_directory(user_id, session_id)