AI Chat Interface

# server/chat/backend/agent/access.py
class ModeAccessController:
    @staticmethod
    def is_tool_allowed(mode: str, tool_name: str) -> bool:
        """Check if a tool is allowed in the current mode."""
        if mode == "ask":
            # Block write operations in ask mode
            write_tools = [
                "cloud_tool", "kubectl_tool", "github_commit",
                "iac_tool", "github_apply_fix"
            ]
            return tool_name not in write_tools
        return True  # Agent mode allows all tools

# server/chat/backend/agent/workflow.py
async def stream(self, input_state: State):
    """Stream the workflow with token-level streaming."""
    async for event in self.app.astream_events(input_state, self.config, version="v2"):
        event_type = event.get("event")
        
        if event_type == "on_chat_model_stream":
            # Stream individual tokens from LLM
            chunk_data = event.get("data", {})
            chunk_obj = chunk_data.get("chunk")
            if chunk_obj and hasattr(chunk_obj, 'content') and chunk_obj.content:
                content = _extract_text_from_content(chunk_obj.content)
                if content:
                    yield ("token", content)

def _consolidate_message_chunks(self):
    """Consolidate AIMessageChunk objects into proper AIMessage objects."""
    chunk_builders = {}
    
    for msg in messages:
        if isinstance(msg, AIMessageChunk):
            builder = chunk_builders.get(msg.id)
            if builder is None:
                builder = {
                    "content": "",
                    "tool_calls": [],
                    "id": msg.id
                }
                chunk_builders[msg.id] = builder
            
            # Accumulate content
            builder["content"] += msg.content
            
            # Accumulate tool calls
            if hasattr(msg, 'tool_calls') and msg.tool_calls:
                builder["tool_calls"].extend(msg.tool_calls)
            
            # Finalize on finish_reason
            if msg.response_metadata.get("finish_reason") in ("tool_calls", "stop"):
                cleaned_tool_calls = self._clean_tool_calls(builder["tool_calls"])
                ai_msg = AIMessage(
                    content=builder["content"],
                    tool_calls=cleaned_tool_calls,
                    id=builder["id"]
                )
                consolidated_messages.append(ai_msg)
                del chunk_builders[msg.id]

// client/src/app/chat/components/useChatSendHandlers.ts
const handleSend = async (text: string, websocket: WebSocket) => {
  let effectiveSessionId = currentSessionId;
  
  if (!effectiveSessionId) {
    // Create new session
    const newSession = await createSession({
      title: text.substring(0, 50) + (text.length > 50 ? '...' : ''),
      userId: userId
    });
    effectiveSessionId = newSession.id;
    setCurrentSessionId(newSession.id);
    justCreatedSessionRef.current = newSession.id;
  }
  
  // Send message via WebSocket
  websocket.send(JSON.stringify({
    type: 'message',
    query: text,
    session_id: effectiveSessionId,
    user_id: userId,
    mode: selectedMode,
    model: selectedModel,
    provider_preference: selectedProviders
  }));
};

# server/chat/backend/agent/workflow.py
if input_state.session_id and input_state.user_id:
    existing_context = LLMContextManager.load_context_history(
        input_state.session_id, 
        input_state.user_id
    )
    
    if existing_context:
        # Combine existing context with new messages
        combined_messages = []
        combined_messages.extend(existing_context)
        combined_messages.extend(input_state.messages)
        input_state.messages = combined_messages

# server/chat/backend/agent/utils/chat_context_manager.py
compressed_messages, was_compressed = ChatContextManager.compress_context_if_needed(
    session_id, user_id, messages_full, model
)

if was_compressed:
    # Use compressed version for LLM context
    messages_for_context = compressed_messages
    logger.info(f"Context compression: {len(compressed_messages)} messages preserved")

// client/src/hooks/useChatCancellation.ts
const cancelCurrentMessage = async () => {
  if (!userId || !sessionId || !webSocket.isConnected) return;
  
  // Send cancel control message
  webSocket.send(JSON.stringify({
    type: 'control',
    action: 'cancel',
    session_id: sessionId,
    user_id: userId
  }));
  
  // Backend cancels the Celery task and consolidates partial messages
};

class RateLimiter:
    def __init__(self, rate, per):
        self.tokens = defaultdict(lambda: rate)
        self.rate = rate
        self.per = per
    
    def is_allowed(self, client_id):
        now = time.time()
        elapsed = now - self.last_checked[client_id]
        self.tokens[client_id] += elapsed * (self.rate / self.per)
        
        if self.tokens[client_id] >= 1:
            self.tokens[client_id] -= 1
            return True
        return False

rate_limiter = RateLimiter(rate=5, per=60)

# After workflow completion
if user_id:
    asyncio.create_task(update_api_cost_cache_async(user_id))
    is_cached, total_cost = get_cached_api_cost(user_id)
    
    # Send usage info to client
    usage_info_response = {
        "type": "usage_info",
        "data": {
            "total_cost": round(total_cost, 2)
        }
    }
    await websocket.send(json.dumps(usage_info_response))