CopilotForXcode/Tool/Sources/LangChain/Chain.swift at feature/github-copilot-proxy · ijunfly/CopilotForXcode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
import Foundation

public protocol Chain {
    associatedtype Input
    associatedtype Output
    func callLogic(_ input: Input, callbackManagers: [ChainCallbackManager]) async throws -> Output
    func parseOutput(_ output: Output) -> String
}

public extension Chain {
    func run(_ input: Input, callbackManagers: [ChainCallbackManager] = []) async throws -> String {
        let output = try await call(input, callbackManagers: callbackManagers)
        return parseOutput(output)
    }

    func call(_ input: Input, callbackManagers: [ChainCallbackManager]) async throws -> Output {
        for callbackManager in callbackManagers {
            callbackManager.onChainStart(type: Self.self, input: input)
        }
        return try await callLogic(input, callbackManagers: callbackManagers)
    }
}

public struct SimpleChain<Input, Output>: Chain {
    let block: (Input) async throws -> Output
    let parseOutputBlock: (Output) -> String

    public init(
        block: @escaping (Input) async throws -> Output,
        parseOutput: @escaping (Output) -> String = { String(describing: $0) }
    ) {
        self.block = block
        parseOutputBlock = parseOutput
    }

    public func callLogic(
        _ input: Input,
        callbackManagers: [ChainCallbackManager]
    ) async throws -> Output {
        return try await block(input)
    }

    public func parseOutput(_ output: Output) -> String {
        return parseOutputBlock(output)
    }
}

public struct ConnectedChain<A: Chain, B: Chain>: Chain where B.Input == A.Output {
    public typealias Input = A.Input
    public typealias Output = (B.Output, A.Output)

    public let chainA: A
    public let chainB: B

    public func callLogic(
        _ input: Input,
        callbackManagers: [ChainCallbackManager] = []
    ) async throws -> Output {
        let a = try await chainA.call(input, callbackManagers: callbackManagers)
        let b = try await chainB.call(a, callbackManagers: callbackManagers)
        return (b, a)
    }

    public func parseOutput(_ output: Output) -> String {
        chainB.parseOutput(output.0)
    }
}

public struct PairedChain<A: Chain, B: Chain>: Chain {
    public typealias Input = (A.Input, B.Input)
    public typealias Output = (A.Output, B.Output)

    public let chainA: A
    public let chainB: B

    public func callLogic(
        _ input: Input,
        callbackManagers: [ChainCallbackManager] = []
    ) async throws -> Output {
        async let a = chainA.call(input.0, callbackManagers: callbackManagers)
        async let b = chainB.call(input.1, callbackManagers: callbackManagers)
        return try await (a, b)
    }

    public func parseOutput(_ output: (A.Output, B.Output)) -> String {
        [chainA.parseOutput(output.0), chainB.parseOutput(output.1)].joined(separator: "\n")
    }
}

public struct MappedChain<A: Chain, NewOutput>: Chain {
    public typealias Input = A.Input
    public typealias Output = NewOutput

    public let chain: A
    public let map: (A.Output) -> NewOutput

    public func callLogic(
        _ input: Input,
        callbackManagers: [ChainCallbackManager]
    ) async throws -> Output {
        let output = try await chain.call(input, callbackManagers: callbackManagers)
        return map(output)
    }

    public func parseOutput(_ output: Output) -> String {
        String(describing: output)
    }
}

public extension Chain {
    func pair<C: Chain>(with another: C) -> PairedChain<Self, C> {
        PairedChain(chainA: self, chainB: another)
    }

    func chain<C: Chain>(to another: C) -> ConnectedChain<Self, C> {
        ConnectedChain(chainA: self, chainB: another)
    }

    func map<NewOutput>(_ map: @escaping (Output) -> NewOutput) -> MappedChain<Self, NewOutput> {
        MappedChain(chain: self, map: map)
    }
}