sleepy_agent_ios/SleepyAgent/Inference/ToolCalling.swift

import Foundation

// MARK: - Tool Protocol

/// Protocol for tools that can be called by the Agent
public protocol Tool: Sendable {
    var name: String { get }
    var displayName: String { get }
    var description: String { get }
    
    func execute(arguments: [String: String]) async throws -> String
}

// MARK: - Web Search Tool

/// Web search tool using SearxNG API
public actor WebSearchTool: Tool {
    public let name = "web_search"
    public let displayName = "Web Search"
    public let description = "Search the web for information. Parameters: query (string, required)"
    
    private var baseUrl: String
    private let urlSession: URLSession
    
    /// Response models for SearxNG API
    private struct SearxngResponse: Codable {
        let query: String
        let results: [SearchResult]
    }
    
    private struct SearchResult: Codable {
        let title: String
        let url: String
        let content: String
        let engine: String?
    }
    
    public init(baseUrl: String, urlSession: URLSession = .shared) {
        self.baseUrl = baseUrl.trimmingCharacters(in: CharacterSet(charactersIn: "/"))
        self.urlSession = urlSession
    }
    
    /// Updates the base URL for the search endpoint
    public func updateBaseUrl(_ newUrl: String) {
        self.baseUrl = newUrl.trimmingCharacters(in: CharacterSet(charactersIn: "/"))
    }
    
    public func execute(arguments: [String: String]) async throws -> String {
        guard let query = arguments["query"] else {
            return "Error: 'query' parameter is required"
        }
        
        return try await performSearch(query: query)
    }
    
    private func performSearch(query: String) async throws -> String {
        guard var urlComponents = URLComponents(string: "\(baseUrl)/search") else {
            throw ToolError.invalidURL
        }
        
        urlComponents.queryItems = [
            URLQueryItem(name: "q", value: query),
            URLQueryItem(name: "format", value: "json"),
            URLQueryItem(name: "safesearch", value: "0")
        ]
        
        guard let url = urlComponents.url else {
            throw ToolError.invalidURL
        }
        
        let (data, response) = try await urlSession.data(from: url)
        
        guard let httpResponse = response as? HTTPURLResponse else {
            throw ToolError.invalidResponse
        }
        
        guard httpResponse.statusCode == 200 else {
            throw ToolError.httpError(statusCode: httpResponse.statusCode)
        }
        
        let decoder = JSONDecoder()
        decoder.keyDecodingStrategy = .convertFromSnakeCase
        
        let searchResponse = try decoder.decode(SearxngResponse.self, from: data)
        
        if searchResponse.results.isEmpty {
            return "No results found for '\(query)'"
        }
        
        return searchResponse.results.prefix(5).map { result in
            """
            Title: \(result.title)
            URL: \(result.url)
            Content: \(result.content)
            """
        }.joined(separator: "\n\n")
    }
}

// MARK: - Tool Errors

public enum ToolError: LocalizedError {
    case invalidURL
    case invalidResponse
    case httpError(statusCode: Int)
    case executionFailed(underlying: Error)
    
    public var errorDescription: String? {
        switch self {
        case .invalidURL:
            return "Invalid URL for tool execution"
        case .invalidResponse:
            return "Invalid response from tool"
        case .httpError(let statusCode):
            return "HTTP error: \(statusCode)"
        case .executionFailed(let error):
            return "Tool execution failed: \(error.localizedDescription)"
        }
    }
}

// MARK: - Tool Result

/// Result of a tool execution
public struct ToolResult: Sendable, Identifiable {
    public let id: String
    public let name: String
    public let result: String
    
    public init(id: String, name: String, result: String) {
        self.id = id
        self.name = name
        self.result = result
    }
}

// MARK: - Tool Parser

/// Parses tool calls from LLM responses
public struct ToolParser: Sendable {
    
    /// Supported tool call patterns
    private let patterns: [(start: String, end: String)] = [
        ("<|tool_call>call:", "<tool_call|>"),
        ("<|tool_call>", "<tool_call|>"),
        ("<tool_call>", "</tool_call>")
    ]
    
    /// Parses tool calls from model response
    /// Handles multiple Gemma tool call formats
    public func parseToolCalls(from response: String) -> [ToolCall] {
        var toolCalls: [ToolCall] = []
        
        for (startTag, endTag) in patterns {
            var currentIndex = response.startIndex
            
            while true {
                guard let startRange = response.range(of: startTag, range: currentIndex..<response.endIndex) else {
                    break
                }
                
                let contentStart = startRange.upperBound
                
                guard let endRange = response.range(of: endTag, range: contentStart..<response.endIndex) else {
                    break
                }
                
                let content = String(response[contentStart..<endRange.lowerBound]).trimmingCharacters(in: .whitespacesAndNewlines)
                
                if let toolCall = parseToolCallContent(content) {
                    toolCalls.append(toolCall)
                }
                
                currentIndex = endRange.upperBound
            }
        }
        
        return toolCalls
    }
    
    /// Extracts content before any tool calls
    public func extractContentBeforeTools(from response: String) -> String {
        var firstIndex: String.Index? = nil
        
        for (startTag, _) in patterns {
            if let range = response.range(of: startTag) {
                if firstIndex == nil || range.lowerBound < firstIndex! {
                    firstIndex = range.lowerBound
                }
            }
        }
        
        if let index = firstIndex {
            return String(response[..<index]).trimmingCharacters(in: .whitespacesAndNewlines)
        }
        
        return response.trimmingCharacters(in: .whitespacesAndNewlines)
    }
    
    /// Parse individual tool call content
    /// Handles:
    /// - tool_name{"name": "...", "arguments": {...}}
    /// - tool_name{query: "value"}
    /// - tool_name{query:<|"|>value<|"|>}
    /// - {"name": "...", ...} (legacy)
    private func parseToolCallContent(_ content: String) -> ToolCall? {
        // Clean up special quote tokens first
        let cleaned = content
            .replacingOccurrences(of: "<|\"|>", with: "\"")
            .replacingOccurrences(of: "\">|>", with: "\"")
        
        // Has braces - parse as {args} or tool_name{args}
        if let braceRange = cleaned.range(of: "{") {
            let toolNamePart = String(cleaned[..<braceRange.lowerBound]).trimmingCharacters(in: .whitespacesAndNewlines)
            let inner = String(cleaned[braceRange.lowerBound...])
            
            // Try JSON first, then direct args
            return parseAsJson(toolNamePrefix: toolNamePart, jsonStr: inner)
                ?? parseAsDirectArgs(toolName: toolNamePart, argsStr: inner)
        }
        
        // Try as pure JSON
        if cleaned.trimmingCharacters(in: .whitespacesAndNewlines).hasPrefix("{") {
            return parseAsJson(toolNamePrefix: "", jsonStr: cleaned)
        }
        
        return nil
    }
    
    private func parseAsJson(toolNamePrefix: String, jsonStr: String) -> ToolCall? {
        guard let data = jsonStr.data(using: .utf8) else {
            return nil
        }
        
        do {
            guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] else {
                return nil
            }
            
            // Get tool name from "name" field or prefix
            let toolName = (json["name"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines)
                ?? (toolNamePrefix.isEmpty ? nil : toolNamePrefix)
            
            guard let name = toolName, !name.isEmpty else {
                return nil
            }
            
            // Get arguments from "arguments" field or top-level
            let args: [String: String]
            if let arguments = json["arguments"] as? [String: Any] {
                args = arguments.compactMapValues { value in
                    if let stringValue = value as? String {
                        return stringValue
                    }
                    return String(describing: value)
                }
            } else {
                args = json
                    .filter { $0.key != "name" }
                    .compactMapValues { value in
                        if let stringValue = value as? String {
                            return stringValue
                        }
                        return String(describing: value)
                    }
            }
            
            return ToolCall(id: generateToolCallId(), name: name, arguments: args)
            
        } catch {
            return nil
        }
    }
    
    private func parseAsDirectArgs(toolName: String, argsStr: String) -> ToolCall? {
        // Extract content between outer braces
        let inner = argsStr.trimmingCharacters(in: .whitespacesAndNewlines)
            .trimmingCharacters(in: CharacterSet(charactersIn: "{}"))
        
        var args: [String: String] = [:]
        var depth = 0
        var current = ""
        var parts: [String] = []
        
        // Split by comma, respecting nested structures
        for char in inner {
            switch char {
            case "{", "[":
                depth += 1
                current.append(char)
            case "}", "]":
                depth -= 1
                current.append(char)
            case ",":
                if depth == 0 {
                    parts.append(current.trimmingCharacters(in: .whitespacesAndNewlines))
                    current = ""
                } else {
                    current.append(char)
                }
            default:
                current.append(char)
            }
        }
        
        if !current.isEmpty {
            parts.append(current.trimmingCharacters(in: .whitespacesAndNewlines))
        }
        
        // Parse each key:value pair
        for part in parts {
            if let colonRange = part.range(of: ":") {
                let key = String(part[..<colonRange.lowerBound])
                    .trimmingCharacters(in: .whitespacesAndNewlines)
                    .trimmingCharacters(in: CharacterSet(charactersIn: "\"'"))
                
                var value = String(part[colonRange.upperBound...])
                    .trimmingCharacters(in: .whitespacesAndNewlines)
                    .trimmingCharacters(in: CharacterSet(charactersIn: "\"'{}"))
                
                if !key.isEmpty {
                    args[key] = value
                }
            }
        }
        
        guard !args.isEmpty && !toolName.isEmpty else {
            return nil
        }
        
        return ToolCall(id: generateToolCallId(), name: toolName, arguments: args)
    }
    
    private func generateToolCallId() -> String {
        let uuid = UUID().uuidString
        let prefix = String(uuid.prefix(8))
        return "call_\(prefix)"
    }
}

// MARK: - Tool Registry

/// Registry for managing available tools
public actor ToolRegistry {
    private var tools: [String: any Tool] = [:]
    
    public init() {}
    
    /// Register a tool
    public func register(_ tool: any Tool) {
        tools[tool.name] = tool
    }
    
    /// Unregister a tool by name
    public func unregister(name: String) {
        tools.removeValue(forKey: name)
    }
    
    /// Get a tool by name
    public func getTool(named name: String) -> (any Tool)? {
        tools[name]
    }
    
    /// Get all registered tools
    public func getAllTools() -> [any Tool] {
        Array(tools.values)
    }
    
    /// Execute a tool with given arguments
    public func executeTool(named name: String, arguments: [String: String]) async -> ToolResult {
        let id = generateToolCallId()
        
        guard let tool = tools[name] else {
            return ToolResult(id: id, name: name, result: "Tool '\(name)' not found")
        }
        
        do {
            let result = try await tool.execute(arguments: arguments)
            return ToolResult(id: id, name: name, result: result)
        } catch {
            return ToolResult(id: id, name: name, result: "Error: \(error.localizedDescription)")
        }
    }
    
    /// Execute multiple tool calls
    public func executeToolCalls(_ toolCalls: [ToolCall]) async -> [ToolResult] {
        var results: [ToolResult] = []
        for toolCall in toolCalls {
            let result = await executeTool(named: toolCall.name, arguments: toolCall.arguments)
            results.append(result)
        }
        return results
    }
    
    /// Build system prompt snippet describing available tools
    public func buildToolsDescription() -> String {
        guard !tools.isEmpty else { return "" }
        
        var description = "\nAvailable tools:\n"
        for (_, tool) in tools.sorted(by: { $0.key < $1.key }) {
            description += "- \(tool.name): \(tool.description)\n"
        }
        
        description += """
        \nWhen you need to use a tool, you MUST output EXACTLY in this format:
        <|tool_call>call:web_search{"name": "web_search", "arguments": {"query": "your search query here"}}<tool_call|>
        
        IMPORTANT:
        - Replace 'web_search' with the actual tool name you want to use
        - Do NOT use 'tool_name' as a placeholder - use the real tool name
        """
        
        return description
    }
    
    private func generateToolCallId() -> String {
        let uuid = UUID().uuidString
        let prefix = String(uuid.prefix(8))
        return "call_\(prefix)"
    }
}