Skip to content

Commit ea120fc

Browse files
TesAntiTesAnti
committed
Hook for chain execution tracing
1 parent 3e12067 commit ea120fc

File tree

6 files changed

+165
-18
lines changed

6 files changed

+165
-18
lines changed

‎LangChain.sln

+3-10
Original file line numberDiff line numberDiff line change
@@ -164,15 +164,13 @@ Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "LangChain.Providers.Google.
164164
EndProject
165165
Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "LangChain.Providers.Azure", "src\libs\Providers\LangChain.Providers.Azure\LangChain.Providers.Azure.csproj", "{18F5AAB1-1750-41BD-B623-6339CA5754D9}"
166166
EndProject
167-
Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "LangChain.Providers.Ollama.IntegrationTests", "src\tests\LangChain.Providers.Ollama.IntegrationTests\LangChain.Providers.Ollama.IntegrationTests.csproj", "{72B1E2CC-1A34-470E-A579-034CB0972BB7}"
167+
Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "LangChain.Providers.Ollama.IntegrationTests", "src\tests\LangChain.Providers.Ollama.IntegrationTests\LangChain.Providers.Ollama.IntegrationTests.csproj", "{72B1E2CC-1A34-470E-A579-034CB0972BB7}"
168168
EndProject
169169
Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "LangChain.Providers.Ollama", "src\libs\Providers\LangChain.Providers.Ollama\LangChain.Providers.Ollama.csproj", "{4913844F-74EC-4E74-AE8A-EA825569E6BA}"
170170
EndProject
171-
Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "LangChain.Providers.Automatic1111", "src\libs\Providers\LangChain.Providers.Automatic1111\LangChain.Providers.Automatic1111.csproj", "{BF4C7B87-0997-4208-84EF-D368DF7B9861}"
171+
Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "LangChain.Providers.Automatic1111", "src\libs\Providers\LangChain.Providers.Automatic1111\LangChain.Providers.Automatic1111.csproj", "{BF4C7B87-0997-4208-84EF-D368DF7B9861}"
172172
EndProject
173-
Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "LangChain.Providers.Automatic1111.IntegrationTests", "src\tests\LangChain.Providers.Automatic1111.IntegrationTests\LangChain.Providers.Automatic1111.IntegrationTests.csproj", "{A6CF79BC-8365-46E8-9230-1A4AD615D40B}"
174-
EndProject
175-
Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "LangChain.Providers.Azure", "src\libs\Providers\LangChain.Providers.Azure\LangChain.Providers.Azure.csproj", "{738984A2-7D3F-42E7-9B4D-3528E2539197}"
173+
Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "LangChain.Providers.Automatic1111.IntegrationTests", "src\tests\LangChain.Providers.Automatic1111.IntegrationTests\LangChain.Providers.Automatic1111.IntegrationTests.csproj", "{A6CF79BC-8365-46E8-9230-1A4AD615D40B}"
176174
EndProject
177175
Global
178176
GlobalSection(SolutionConfigurationPlatforms) = preSolution
@@ -404,10 +402,6 @@ Global
404402
{A6CF79BC-8365-46E8-9230-1A4AD615D40B}.Debug|Any CPU.Build.0 = Debug|Any CPU
405403
{A6CF79BC-8365-46E8-9230-1A4AD615D40B}.Release|Any CPU.ActiveCfg = Release|Any CPU
406404
{A6CF79BC-8365-46E8-9230-1A4AD615D40B}.Release|Any CPU.Build.0 = Release|Any CPU
407-
{738984A2-7D3F-42E7-9B4D-3528E2539197}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
408-
{738984A2-7D3F-42E7-9B4D-3528E2539197}.Debug|Any CPU.Build.0 = Debug|Any CPU
409-
{738984A2-7D3F-42E7-9B4D-3528E2539197}.Release|Any CPU.ActiveCfg = Release|Any CPU
410-
{738984A2-7D3F-42E7-9B4D-3528E2539197}.Release|Any CPU.Build.0 = Release|Any CPU
411405
EndGlobalSection
412406
GlobalSection(SolutionProperties) = preSolution
413407
HideSolutionNode = FALSE
@@ -475,7 +469,6 @@ Global
475469
{4913844F-74EC-4E74-AE8A-EA825569E6BA} = {E55391DE-F8F3-4CC2-A0E3-2406C76E9C68}
476470
{BF4C7B87-0997-4208-84EF-D368DF7B9861} = {E55391DE-F8F3-4CC2-A0E3-2406C76E9C68}
477471
{A6CF79BC-8365-46E8-9230-1A4AD615D40B} = {FDEE2E22-C239-4921-83B2-9797F765FD6A}
478-
{738984A2-7D3F-42E7-9B4D-3528E2539197} = {E55391DE-F8F3-4CC2-A0E3-2406C76E9C68}
479472
EndGlobalSection
480473
GlobalSection(ExtensibilityGlobals) = postSolution
481474
SolutionGuid = {5C00D0F1-6138-4ED9-846B-97E43D6DFF1C}

‎src/libs/LangChain.Core/Chains/StackableChains/BaseStackableChain.cs

+19-7
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
using LangChain.Abstractions.Schema;
33
using LangChain.Callback;
44
using LangChain.Chains.HelperChains.Exceptions;
5+
using LangChain.Chains.StackableChains.Context;
56
using LangChain.Schema;
67

78
namespace LangChain.Chains.HelperChains;
@@ -86,10 +87,13 @@ string FormatInputValues(IChainValues values)
8687
public Task<IChainValues> CallAsync(IChainValues values, ICallbacks? callbacks = null,
8788
IReadOnlyList<string>? tags = null, IReadOnlyDictionary<string, object>? metadata = null)
8889
{
90+
91+
8992
if (values == null)
9093
{
9194
throw new ArgumentNullException(nameof(values));
9295
}
96+
9397
try
9498
{
9599
return InternalCall(values);
@@ -108,8 +112,9 @@ public Task<IChainValues> CallAsync(IChainValues values, ICallbacks? callbacks =
108112

109113
throw new StackableChainException(message, ex);
110114
}
115+
111116
}
112-
117+
113118
/// <summary>
114119
///
115120
/// </summary>
@@ -143,9 +148,11 @@ public static StackChain BitwiseOr(BaseStackableChain left, BaseStackableChain r
143148
///
144149
/// </summary>
145150
/// <returns></returns>
146-
public async Task<IChainValues> Run()
151+
public async Task<IChainValues> Run(StackableChainHook? hook=null)
147152
{
148-
var res = await CallAsync(new ChainValues()).ConfigureAwait(false);
153+
var values = new StackableChainValues() {Hook = hook};
154+
hook?.ChainStart(values);
155+
var res = await CallAsync(values).ConfigureAwait(false);
149156
return res;
150157
}
151158

@@ -154,9 +161,9 @@ public async Task<IChainValues> Run()
154161
/// </summary>
155162
/// <param name="resultKey"></param>
156163
/// <returns></returns>
157-
public async Task<string?> Run(string resultKey)
164+
public async Task<string?> Run(string resultKey, StackableChainHook? hook = null)
158165
{
159-
var res = await CallAsync(new ChainValues()).ConfigureAwait(false);
166+
var res = await CallAsync(new StackableChainValues() { Hook = hook }).ConfigureAwait(false);
160167
return res.Value[resultKey].ToString();
161168
}
162169

@@ -166,12 +173,17 @@ public async Task<IChainValues> Run()
166173
/// <param name="resultKey"></param>
167174
/// <typeparam name="T"></typeparam>
168175
/// <returns></returns>
169-
public async Task<T> Run<T>(string resultKey)
176+
public async Task<T> Run<T>(string resultKey, StackableChainHook? hook = null)
170177
{
171-
var res = await CallAsync(new ChainValues()).ConfigureAwait(false);
178+
var res = await CallAsync(new StackableChainValues() { Hook = hook }).ConfigureAwait(false);
172179
return (T)res.Value[resultKey];
173180
}
174181

182+
public Task<string?> Run(string resultKey)
183+
{
184+
return Run(resultKey, null);
185+
}
186+
175187
/// <summary>
176188
///
177189
/// </summary>

‎src/libs/LangChain.Core/Chains/StackableChains/Hooks/ConsoleTraceHook.cs

+102
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+

2+
using LangChain.Chains.HelperChains;
3+
4+
namespace LangChain.Chains.StackableChains.Context;
5+
6+
public class ConsoleTraceHook: StackableChainHook
7+
{
8+
public bool UseColors { get; set; }=true;
9+
public int ValuesLength { get; set; } = 40;
10+
public override void ChainStart(StackableChainValues values)
11+
{
12+
Console.WriteLine();
13+
}
14+
public override void LinkEnter(BaseStackableChain chain, StackableChainValues values)
15+
{
16+
17+
Console.Write("|");
18+
Console.Write(chain.GetType().Name);
19+
Console.WriteLine();
20+
if (chain.InputKeys.Count > 0)
21+
{
22+
Console.Write("|");
23+
Console.Write(" ");
24+
Console.Write("\u2514");
25+
Console.Write("Input:");
26+
Console.WriteLine();
27+
foreach (string inputKey in chain.InputKeys)
28+
{
29+
Console.Write("|");
30+
Console.Write(" ");
31+
Console.Write("\u2514");
32+
var value = values.Value[inputKey];
33+
var oldColor = Console.ForegroundColor;
34+
Console.ForegroundColor = GetColorForKey(inputKey);
35+
Console.Write($" {inputKey}={ShortenString(value.ToString() ?? "", ValuesLength)}");
36+
Console.ForegroundColor = oldColor;
37+
Console.WriteLine();
38+
}
39+
}
40+
41+
42+
}
43+
44+
Dictionary<string, ConsoleColor> _colorMap = new Dictionary<string, ConsoleColor>();
45+
46+
ConsoleColor GetColorForKey(string key)
47+
{
48+
if(!UseColors)
49+
return Console.ForegroundColor;
50+
// if key is not in map, get unique color(except black and white)
51+
// if there no unique colors left, return white
52+
if (!_colorMap.ContainsKey(key))
53+
{
54+
var color = ConsoleColor.White;
55+
var colors = Enum.GetValues(typeof(ConsoleColor));
56+
foreach (ConsoleColor c in colors)
57+
{
58+
if (c == ConsoleColor.Black || c == ConsoleColor.White)
59+
continue;
60+
if (!_colorMap.ContainsValue(c))
61+
{
62+
color = c;
63+
break;
64+
}
65+
}
66+
_colorMap.Add(key, color);
67+
}
68+
return _colorMap[key];
69+
}
70+
71+
string ShortenString(string str, int length)
72+
{
73+
if (str.Length <= length)
74+
return str;
75+
return str.Substring(0, length - 3) + "...";
76+
}
77+
78+
public override void LinkExit(BaseStackableChain chain, StackableChainValues values)
79+
{
80+
if (chain.OutputKeys.Count > 0)
81+
{
82+
Console.Write("|");
83+
Console.Write(" ");
84+
Console.Write("\u2514");
85+
Console.Write("Output:");
86+
Console.WriteLine();
87+
foreach (string outputKey in chain.OutputKeys)
88+
{
89+
Console.Write("|");
90+
Console.Write(" ");
91+
Console.Write("\u2514");
92+
var value = values.Value[outputKey];
93+
var oldColor = Console.ForegroundColor;
94+
Console.ForegroundColor = GetColorForKey(outputKey);
95+
Console.Write($" {outputKey}={ShortenString(value.ToString() ?? "", ValuesLength)}");
96+
Console.ForegroundColor = oldColor;
97+
Console.WriteLine();
98+
}
99+
}
100+
101+
}
102+
}

‎src/libs/LangChain.Core/Chains/StackableChains/Hooks/StackableChainHook.cs

+21
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
using LangChain.Chains.HelperChains;
2+
3+
namespace LangChain.Chains.StackableChains.Context;
4+
5+
public class StackableChainHook
6+
{
7+
public virtual void ChainStart(StackableChainValues values)
8+
{
9+
10+
}
11+
12+
public virtual void LinkEnter(BaseStackableChain chain, StackableChainValues values)
13+
{
14+
15+
}
16+
17+
public virtual void LinkExit(BaseStackableChain chain, StackableChainValues values)
18+
{
19+
20+
}
21+
}

‎src/libs/LangChain.Core/Chains/StackableChains/Hooks/StackableChainValues.cs

+8
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
using LangChain.Schema;
2+
3+
namespace LangChain.Chains.StackableChains.Context;
4+
5+
public class StackableChainValues : ChainValues
6+
{
7+
public StackableChainHook? Hook { get; set; }
8+
}

‎src/libs/LangChain.Core/Chains/StackableChains/StackChain.cs

+12-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
using LangChain.Abstractions.Schema;
2+
using LangChain.Chains.StackableChains.Context;
23
using LangChain.Schema;
34

45
namespace LangChain.Chains.HelperChains;
@@ -60,15 +61,25 @@ protected override async Task<IChainValues> InternalCall(IChainValues values)
6061

6162
if (IsolatedInputKeys.Count > 0)
6263
{
63-
var res = new ChainValues();
64+
var res = new StackableChainValues(){Hook = (values as StackableChainValues)?.Hook};
6465
foreach (var key in IsolatedInputKeys)
6566
{
6667
res.Value[key] = values.Value[key];
6768
}
6869
values = res;
6970
}
71+
if(a is not StackChain)
72+
(values as StackableChainValues)?.Hook?.LinkEnter(a, (values as StackableChainValues));
7073
await a.CallAsync(values).ConfigureAwait(false);
74+
if (a is not StackChain)
75+
(values as StackableChainValues)?.Hook?.LinkExit(a, (values as StackableChainValues));
76+
77+
if (b is not StackChain)
78+
(values as StackableChainValues)?.Hook?.LinkEnter(b, (values as StackableChainValues));
7179
await b.CallAsync(values).ConfigureAwait(false);
80+
if (b is not StackChain)
81+
(values as StackableChainValues)?.Hook?.LinkExit(b, (values as StackableChainValues));
82+
7283
if (IsolatedOutputKeys.Count > 0)
7384
{
7485
foreach (var key in IsolatedOutputKeys)

0 commit comments

Comments
 (0)