Skip to content

Commit bd6cc94

Browse files
dannykoppingdannykopping
committed
fixup tool calling
Signed-off-by: Danny Kopping <dannykopping@gmail.com>
1 parent e4cd696 commit bd6cc94

File tree

3 files changed

+10
-36
lines changed

3 files changed

+10
-36
lines changed

‎aibridged/bridge.go

Lines changed: 5 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -494,7 +494,8 @@ func (b *Bridge) proxyOpenAIRequest(w http.ResponseWriter, r *http.Request) {
494494
appendedPrevMsg := false
495495
// Process each pending tool call
496496
for _, tc := range pendingToolCalls {
497-
if !b.isInjectedTool(tc.ID) {
497+
fn := tc.Function.Name
498+
if !b.isInjectedTool(fn) {
498499
// Not an MCP proxy call, don't do anything.
499500
continue
500501
}
@@ -508,7 +509,7 @@ func (b *Bridge) proxyOpenAIRequest(w http.ResponseWriter, r *http.Request) {
508509
_, err = coderdClient.TrackToolUsage(ctx, &proto.TrackToolUsageRequest{
509510
Model: string(in.Model),
510511
Input: tc.Function.Arguments,
511-
Tool: tc.Function.Name,
512+
Tool: fn,
512513
Injected: true,
513514
})
514515
if err != nil {
@@ -521,7 +522,7 @@ func (b *Bridge) proxyOpenAIRequest(w http.ResponseWriter, r *http.Request) {
521522
)
522523
_ = json.NewEncoder(&buf).Encode(tc.Function.Arguments)
523524
_ = json.NewDecoder(&buf).Decode(&args)
524-
res, err := b.tools[tc.ID].Call(ctx, args)
525+
res, err := b.tools[fn].Call(ctx, args)
525526
if err != nil {
526527
// Always provide a tool result even if the tool call failed
527528
errorResponse := map[string]interface{}{
@@ -759,6 +760,7 @@ func (b *Bridge) proxyAnthropicRequest(w http.ResponseWriter, r *http.Request) {
759760
break
760761
}
761762

763+
// TODO: implement injected tool calling.
762764
if resp.StopReason == anthropic.BetaStopReasonToolUse {
763765
var (
764766
toolUse anthropic.BetaToolUseBlock
@@ -809,12 +811,6 @@ func (b *Bridge) proxyAnthropicRequest(w http.ResponseWriter, r *http.Request) {
809811

810812
es := newEventStream(anthropicEventStream)
811813

812-
// var buf strings.Builder
813-
// in.Messages[0].Content = []anthropic.BetaContentBlockParamUnion{in.Messages[0].Content[len(in.Messages[0].Content) - 1]}
814-
//
815-
// json.NewEncoder(&buf).Encode(in)
816-
// fmt.Println(strings.Replace(buf.String(), "'", "\\'", -1))
817-
818814
var wg sync.WaitGroup
819815
wg.Add(1)
820816

@@ -862,8 +858,6 @@ func (b *Bridge) proxyAnthropicRequest(w http.ResponseWriter, r *http.Request) {
862858
// Don't relay this event back, otherwise the client will try invoke the tool as well.
863859
continue
864860
}
865-
default:
866-
fmt.Printf("[%s] %s\n", event.Type, event.RawJSON())
867861
}
868862
case string(constant.ValueOf[ant_constant.ContentBlockDelta]()):
869863
if len(pendingToolCalls) > 0 && b.isInjectedTool(lastToolName) {
@@ -1179,19 +1173,6 @@ func (b *Bridge) proxyAnthropicRequest(w http.ResponseWriter, r *http.Request) {
11791173
}
11801174

11811175
<-streamCtx.Done()
1182-
1183-
// TODO: do we need to do this?
1184-
// // Close the underlying connection by hijacking it
1185-
// if hijacker, ok := w.(http.Hijacker); ok {
1186-
// conn, _, err := hijacker.Hijack()
1187-
// if err != nil {
1188-
// b.logger.Error(ctx, "failed to hijack connection", slog.Error(err))
1189-
// } else {
1190-
// conn.Close() // This closes the TCP connection entirely
1191-
// b.logger.Debug(ctx, "connection closed, stream over")
1192-
// }
1193-
// }
1194-
11951176
break
11961177
}
11971178
}

‎aibridged/mcp.go

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,10 @@ func (b *MCPToolBridge) CallTool(ctx context.Context, name string, input any) (*
9696
}
9797

9898
func (t *MCPTool) Call(ctx context.Context, input any) (*mcp.CallToolResult, error) {
99+
if t == nil {
100+
return nil, xerrors.Errorf("nil tool!")
101+
}
102+
99103
return t.client.CallTool(ctx, mcp.CallToolRequest{
100104
Params: mcp.CallToolParams{
101105
Name: t.Name,
@@ -119,7 +123,7 @@ func (b *MCPToolBridge) fetchMCPTools(ctx context.Context) (map[string]*MCPTool,
119123
if err != nil {
120124
return nil, xerrors.Errorf("init MCP client: %w", err)
121125
}
122-
fmt.Printf("mcp(%q)], %+v\n", result.ServerInfo.Name, result) // TODO: remove.
126+
b.logger.Debug(ctx, "mcp client initialized", slog.F("name", result.ServerInfo.Name), slog.F("server_version", result.ServerInfo.Version))
123127

124128
// Test tool listing
125129
tools, err := b.client.ListTools(ctx, mcp.ListToolsRequest{})

‎coderd/aibridge.go

Lines changed: 0 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -62,17 +62,6 @@ func (api *API) bridgeAIRequest(rw http.ResponseWriter, r *http.Request) {
6262
http.Error(rw, "failed to create ai bridge", http.StatusInternalServerError)
6363
return
6464
}
65-
66-
// TODO: direct func calls instead of HTTP srv?
67-
done := make(chan any, 1)
68-
go func() {
69-
defer close(done)
70-
err := bridge.Serve()
71-
// TODO: better error handling.
72-
// TODO: close on shutdown.
73-
api.Logger.Warn(ctx, "bridge server stopped", slog.Error(err))
74-
}()
75-
7665
http.StripPrefix("/api/v2/aibridge", bridge.Handler()).ServeHTTP(rw, r)
7766
}
7867

0 commit comments

Comments
 (0)