diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..66afed6 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +example/aibridge.db diff --git a/README.md b/README.md index 89ef769..9308b25 100644 --- a/README.md +++ b/README.md @@ -1 +1,131 @@ -aibridge +# aibridge + +aibridge is an HTTP gateway that sits between AI clients and upstream AI providers (Anthropic, OpenAI). It intercepts requests to record token usage, prompts, and tool invocations per user. Optionally supports centralized [MCP](https://modelcontextprotocol.io/) tool injection with allowlist/denylist filtering. + +## Architecture + +``` +┌─────────────────┐ ┌───────────────────────────────────────────┐ +│ AI Client │ │ aibridge │ +│ (Claude Code, │────▶│ ┌─────────────────┐ ┌─────────────┐ │ +│ Cursor, etc.) │ │ │ RequestBridge │───▶│ Providers │ │ +└─────────────────┘ │ │ (http.Handler) │ │ (Anthropic │ │ + │ └─────────────────┘ │ OpenAI) │ │ + │ └──────┬──────┘ │ + │ │ │ + │ ▼ │ ┌─────────────┐ + │ ┌─────────────────┐ ┌─────────────┐ │ │ Upstream │ + │ │ Recorder │◀───│ Interceptor │─── ───▶│ API │ + │ │ (tokens, tools, │ │ (streaming/ │ │ │ (Anthropic │ + │ │ prompts) │ │ blocking) │ │ │ OpenAI) │ + │ └────────┬────────┘ └──────┬──────┘ │ └─────────────┘ + │ │ │ │ + │ ▼ ┌──────▼──────┐ │ + │ ┌ ─ ─ ─ ─ ─ ─ ─ ┐ │ MCP Proxy │ │ + │ │ Database │ │ (tools) │ │ + │ └ ─ ─ ─ ─ ─ ─ ─ ┘ └─────────────┘ │ + └───────────────────────────────────────────┘ +``` + +### Components + +- **RequestBridge**: The main `http.Handler` that routes requests to providers +- **Provider**: Defines bridged routes (intercepted) and passthrough routes (proxied) +- **Interceptor**: Handles request/response processing and streaming +- **Recorder**: Interface for capturing usage data (tokens, prompts, tools) +- **MCP Proxy** (optional): Connects to MCP servers to list tool, inject them into requests, and invoke them in an inner agentic loop + +## Request Flow + +1. Client sends request to `/anthropic/v1/messages` or `/openai/v1/chat/completions` +2. **Actor extraction**: Request must have an actor in context (via `AsActor()`). +3. **Upstream call**: Request forwarded to the AI provider +4. **Response relay**: Response streamed/sent to client +5. **Recording**: Token usage, prompts, and tool invocations recorded + +**With MCP enabled**: Tools from configured MCP servers are centrally defined and injected into requests (prefixed `bmcp_`). Allowlist/denylist regex patterns control which tools are available. When the model selects an injected tool, the gateway invokes it in an inner agentic loop, and continues the conversation loop until complete. + +Passthrough routes (`/v1/models`, `/v1/messages/count_tokens`) are reverse-proxied directly. + +## Observability + +### Prometheus Metrics + +Create metrics with `NewMetrics(prometheus.Registerer)`: + +| Metric | Type | Description | +|--------|------|-------------| +| `interceptions_total` | Counter | Intercepted request count | +| `interceptions_inflight` | Gauge | Currently processing requests | +| `interceptions_duration_seconds` | Histogram | Request duration | +| `tokens_total` | Counter | Token usage (input/output) | +| `prompts_total` | Counter | User prompt count | +| `injected_tool_invocations_total` | Counter | MCP tool invocations | +| `passthrough_total` | Counter | Non-intercepted requests | + +### Recorder Interface + +Implement `Recorder` to persist usage data to your database. The example uses SQLite ([example/recorder.go](example/recorder.go)): + +- `aibridge_interceptions` - request metadata (provider, model, initiator, timestamps) +- `aibridge_token_usages` - input/output token counts per response +- `aibridge_user_prompts` - user prompts +- `aibridge_tool_usages` - tool invocations (injected and client-defined) + +```go +type Recorder interface { + RecordInterception(ctx context.Context, req *InterceptionRecord) error + RecordInterceptionEnded(ctx context.Context, req *InterceptionRecordEnded) error + RecordTokenUsage(ctx context.Context, req *TokenUsageRecord) error + RecordPromptUsage(ctx context.Context, req *PromptUsageRecord) error + RecordToolUsage(ctx context.Context, req *ToolUsageRecord) error +} +``` + +## Example + +See [example/](example/) for a complete runnable example with SQLite persistence and [DeepWiki](https://mcp.deepwiki.com) MCP integration. + +### Setup + +1. **Get API keys** from the provider consoles: + - Anthropic: https://console.anthropic.com/settings/keys + - OpenAI: https://platform.openai.com/api-keys + +2. **Set environment variables**: + ```bash + export ANTHROPIC_API_KEY="sk-ant-..." + export OPENAI_API_KEY="sk-..." + ``` + +3. **Run the example**: + ```bash + cd example && go run . + ``` + +4. **Test with curl**: + ```bash + curl -X POST http://localhost:8080/anthropic/v1/messages \ + -H "Content-Type: application/json" \ + -d '{ + "model": "claude-sonnet-4-20250514", + "max_tokens": 1024, + "messages": [{"role": "user", "content": "Hello!"}], + "stream": true + }' + ``` + +5. **Test with Claude Code**: + Claude Code allows a base URL override via `ANTHROPIC_BASE_URL`. + + ![image with cloude code example](example/claude-code.png) + +## Supported Routes + +| Provider | Route | Type | +|----------|-------|------| +| Anthropic | `/anthropic/v1/messages` | Bridged (intercepted) | +| Anthropic | `/anthropic/v1/models` | Passthrough | +| Anthropic | `/anthropic/v1/messages/count_tokens` | Passthrough | +| OpenAI | `/openai/v1/chat/completions` | Bridged (intercepted) | +| OpenAI | `/openai/v1/models` | Passthrough | diff --git a/example/claude-code.png b/example/claude-code.png new file mode 100644 index 0000000..16bf1aa Binary files /dev/null and b/example/claude-code.png differ diff --git a/example/go.mod b/example/go.mod new file mode 100644 index 0000000..6caf6cf --- /dev/null +++ b/example/go.mod @@ -0,0 +1,77 @@ +module github.com/coder/aibridge/example + +go 1.24.6 + +toolchain go1.24.10 + +require ( + cdr.dev/slog v1.6.2-0.20250703074222-9df5e0a6c145 + github.com/coder/aibridge v0.0.0 + github.com/google/uuid v1.6.0 + github.com/prometheus/client_golang v1.23.2 + modernc.org/sqlite v1.40.1 +) + +require ( + github.com/anthropics/anthropic-sdk-go v1.13.0 // indirect + github.com/aws/aws-sdk-go-v2 v1.30.3 // indirect + github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.3 // indirect + github.com/aws/aws-sdk-go-v2/config v1.27.27 // indirect + github.com/aws/aws-sdk-go-v2/credentials v1.17.27 // indirect + github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.11 // indirect + github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.15 // indirect + github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.15 // indirect + github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.3 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.17 // indirect + github.com/aws/aws-sdk-go-v2/service/sso v1.22.4 // indirect + github.com/aws/aws-sdk-go-v2/service/ssooidc v1.26.4 // indirect + github.com/aws/aws-sdk-go-v2/service/sts v1.30.3 // indirect + github.com/aws/smithy-go v1.20.3 // indirect + github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect + github.com/bahlo/generic-list-go v0.2.0 // indirect + github.com/beorn7/perks v1.0.1 // indirect + github.com/buger/jsonparser v1.1.1 // indirect + github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/charmbracelet/lipgloss v0.7.1 // indirect + github.com/dustin/go-humanize v1.0.1 // indirect + github.com/hashicorp/errwrap v1.0.0 // indirect + github.com/hashicorp/go-multierror v1.1.1 // indirect + github.com/invopop/jsonschema v0.13.0 // indirect + github.com/lucasb-eyer/go-colorful v1.2.0 // indirect + github.com/mailru/easyjson v0.7.7 // indirect + github.com/mark3labs/mcp-go v0.38.0 // indirect + github.com/mattn/go-isatty v0.0.20 // indirect + github.com/mattn/go-runewidth v0.0.15 // indirect + github.com/muesli/reflow v0.3.0 // indirect + github.com/muesli/termenv v0.15.2 // indirect + github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect + github.com/ncruces/go-strftime v0.1.9 // indirect + github.com/openai/openai-go/v2 v2.7.0 // indirect + github.com/prometheus/client_model v0.6.2 // indirect + github.com/prometheus/common v0.66.1 // indirect + github.com/prometheus/procfs v0.16.1 // indirect + github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect + github.com/rivo/uniseg v0.4.4 // indirect + github.com/spf13/cast v1.7.1 // indirect + github.com/tidwall/gjson v1.18.0 // indirect + github.com/tidwall/match v1.2.0 // indirect + github.com/tidwall/pretty v1.2.1 // indirect + github.com/tidwall/sjson v1.2.5 // indirect + github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect + github.com/yosida95/uritemplate/v3 v3.0.2 // indirect + go.opentelemetry.io/otel v1.33.0 // indirect + go.opentelemetry.io/otel/trace v1.33.0 // indirect + go.yaml.in/yaml/v2 v2.4.2 // indirect + golang.org/x/exp v0.0.0-20250819193227-8b4c13bb791b // indirect + golang.org/x/sys v0.36.0 // indirect + golang.org/x/term v0.34.0 // indirect + golang.org/x/xerrors v0.0.0-20240716161551-93cc26a95ae9 // indirect + google.golang.org/protobuf v1.36.8 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect + modernc.org/libc v1.66.10 // indirect + modernc.org/mathutil v1.7.1 // indirect + modernc.org/memory v1.11.0 // indirect +) + +replace github.com/coder/aibridge => ../ diff --git a/example/go.sum b/example/go.sum new file mode 100644 index 0000000..cfc0a61 --- /dev/null +++ b/example/go.sum @@ -0,0 +1,214 @@ +cdr.dev/slog v1.6.2-0.20250703074222-9df5e0a6c145 h1:Mk4axSLxKw3hjkf3PffBLQYta7nPVIWObuKCPDWgQLc= +cdr.dev/slog v1.6.2-0.20250703074222-9df5e0a6c145/go.mod h1:NaoTA7KwopCrnaSb0JXTC0PTp/O/Y83Lndnq0OEV3ZQ= +cloud.google.com/go/compute v1.23.0 h1:tP41Zoavr8ptEqaW6j+LQOnyBBhO7OkOMAGrgLopTwY= +cloud.google.com/go/compute/metadata v0.5.0 h1:Zr0eK8JbFv6+Wi4ilXAR8FJ3wyNdpxHKJNPos6LTZOY= +cloud.google.com/go/compute/metadata v0.5.0/go.mod h1:aHnloV2TPI38yx4s9+wAZhHykWvVCfu7hQbF+9CWoiY= +cloud.google.com/go/logging v1.8.1 h1:26skQWPeYhvIasWKm48+Eq7oUqdcdbwsCVwz5Ys0FvU= +cloud.google.com/go/logging v1.8.1/go.mod h1:TJjR+SimHwuC8MZ9cjByQulAMgni+RkXeI3wwctHJEI= +cloud.google.com/go/longrunning v0.5.1 h1:Fr7TXftcqTudoyRJa113hyaqlGdiBQkp0Gq7tErFDWI= +cloud.google.com/go/longrunning v0.5.1/go.mod h1:spvimkwdz6SPWKEt/XBij79E9fiTkHSQl/fRUUQJYJc= +github.com/anthropics/anthropic-sdk-go v1.13.0 h1:Bhbe8sRoDPtipttg8bQYrMCKe2b79+q6rFW1vOKEUKI= +github.com/anthropics/anthropic-sdk-go v1.13.0/go.mod h1:WTz31rIUHUHqai2UslPpw5CwXrQP3geYBioRV4WOLvE= +github.com/aws/aws-sdk-go-v2 v1.30.3 h1:jUeBtG0Ih+ZIFH0F4UkmL9w3cSpaMv9tYYDbzILP8dY= +github.com/aws/aws-sdk-go-v2 v1.30.3/go.mod h1:nIQjQVp5sfpQcTc9mPSr1B0PaWK5ByX9MOoDadSN4lc= +github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.3 h1:tW1/Rkad38LA15X4UQtjXZXNKsCgkshC3EbmcUmghTg= +github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.3/go.mod h1:UbnqO+zjqk3uIt9yCACHJ9IVNhyhOCnYk8yA19SAWrM= +github.com/aws/aws-sdk-go-v2/config v1.27.27 h1:HdqgGt1OAP0HkEDDShEl0oSYa9ZZBSOmKpdpsDMdO90= +github.com/aws/aws-sdk-go-v2/config v1.27.27/go.mod h1:MVYamCg76dFNINkZFu4n4RjDixhVr51HLj4ErWzrVwg= +github.com/aws/aws-sdk-go-v2/credentials v1.17.27 h1:2raNba6gr2IfA0eqqiP2XiQ0UVOpGPgDSi0I9iAP+UI= +github.com/aws/aws-sdk-go-v2/credentials v1.17.27/go.mod h1:gniiwbGahQByxan6YjQUMcW4Aov6bLC3m+evgcoN4r4= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.11 h1:KreluoV8FZDEtI6Co2xuNk/UqI9iwMrOx/87PBNIKqw= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.11/go.mod h1:SeSUYBLsMYFoRvHE0Tjvn7kbxaUhl75CJi1sbfhMxkU= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.15 h1:SoNJ4RlFEQEbtDcCEt+QG56MY4fm4W8rYirAmq+/DdU= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.15/go.mod h1:U9ke74k1n2bf+RIgoX1SXFed1HLs51OgUSs+Ph0KJP8= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.15 h1:C6WHdGnTDIYETAm5iErQUiVNsclNx9qbJVPIt03B6bI= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.15/go.mod h1:ZQLZqhcu+JhSrA9/NXRm8SkDvsycE+JkV3WGY41e+IM= +github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 h1:hT8rVHwugYE2lEfdFE0QWVo81lF7jMrYJVDWI+f+VxU= +github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0/go.mod h1:8tu/lYfQfFe6IGnaOdrpVgEL2IrrDOf6/m9RQum4NkY= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.3 h1:dT3MqvGhSoaIhRseqw2I0yH81l7wiR2vjs57O51EAm8= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.3/go.mod h1:GlAeCkHwugxdHaueRr4nhPuY+WW+gR8UjlcqzPr1SPI= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.17 h1:HGErhhrxZlQ044RiM+WdoZxp0p+EGM62y3L6pwA4olE= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.17/go.mod h1:RkZEx4l0EHYDJpWppMJ3nD9wZJAa8/0lq9aVC+r2UII= +github.com/aws/aws-sdk-go-v2/service/sso v1.22.4 h1:BXx0ZIxvrJdSgSvKTZ+yRBeSqqgPM89VPlulEcl37tM= +github.com/aws/aws-sdk-go-v2/service/sso v1.22.4/go.mod h1:ooyCOXjvJEsUw7x+ZDHeISPMhtwI3ZCB7ggFMcFfWLU= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.26.4 h1:yiwVzJW2ZxZTurVbYWA7QOrAaCYQR72t0wrSBfoesUE= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.26.4/go.mod h1:0oxfLkpz3rQ/CHlx5hB7H69YUpFiI1tql6Q6Ne+1bCw= +github.com/aws/aws-sdk-go-v2/service/sts v1.30.3 h1:ZsDKRLXGWHk8WdtyYMoGNO7bTudrvuKpDKgMVRlepGE= +github.com/aws/aws-sdk-go-v2/service/sts v1.30.3/go.mod h1:zwySh8fpFyXp9yOr/KVzxOl8SRqgf/IDw5aUt9UKFcQ= +github.com/aws/smithy-go v1.20.3 h1:ryHwveWzPV5BIof6fyDvor6V3iUL7nTfiTKXHiW05nE= +github.com/aws/smithy-go v1.20.3/go.mod h1:krry+ya/rV9RDcV/Q16kpu6ypI4K2czasz0NC3qS14E= +github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k= +github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8= +github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= +github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= +github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= +github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= +github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs= +github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= +github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/charmbracelet/lipgloss v0.7.1 h1:17WMwi7N1b1rVWOjMT+rCh7sQkvDU75B2hbZpc5Kc1E= +github.com/charmbracelet/lipgloss v0.7.1/go.mod h1:yG0k3giv8Qj8edTCbbg6AlQ5e8KNWpFujkNawKNhE2c= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= +github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= +github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8= +github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0= +github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY= +github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= +github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= +github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs= +github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/hashicorp/errwrap v1.0.0 h1:hLrqtEDnRye3+sgx6z4qVLNuviH3MR5aQ0ykNJa/UYA= +github.com/hashicorp/errwrap v1.0.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4= +github.com/hashicorp/go-multierror v1.1.1 h1:H5DkEtf6CXdFp0N0Em5UCwQpXMWke8IA0+lD48awMYo= +github.com/hashicorp/go-multierror v1.1.1/go.mod h1:iw975J/qwKPdAO1clOe2L8331t/9/fmwbPZ6JB6eMoM= +github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E= +github.com/invopop/jsonschema v0.13.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0= +github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= +github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo= +github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= +github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= +github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69Aj6K7nkY= +github.com/lucasb-eyer/go-colorful v1.2.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= +github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= +github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= +github.com/mark3labs/mcp-go v0.38.0 h1:E5tmJiIXkhwlV0pLAwAT0O5ZjUZSISE/2Jxg+6vpq4I= +github.com/mark3labs/mcp-go v0.38.0/go.mod h1:T7tUa2jO6MavG+3P25Oy/jR7iCeJPHImCZHRymCn39g= +github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= +github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/mattn/go-runewidth v0.0.12/go.mod h1:RAqKPSqVFrSLVXbA8x7dzmKdmGzieGRCM46jaSJTDAk= +github.com/mattn/go-runewidth v0.0.15 h1:UNAjwbU9l54TA3KzvqLGxwWjHmMgBUVhBiTjelZgg3U= +github.com/mattn/go-runewidth v0.0.15/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= +github.com/muesli/reflow v0.3.0 h1:IFsN6K9NfGtjeggFP+68I4chLZV2yIKsXJFNZ+eWh6s= +github.com/muesli/reflow v0.3.0/go.mod h1:pbwTDkVPibjO2kyvBQRBxTWEEGDGq0FlB1BIKtnHY/8= +github.com/muesli/termenv v0.15.2 h1:GohcuySI0QmI3wN8Ok9PtKGkgkFIk7y6Vpb5PvrY+Wo= +github.com/muesli/termenv v0.15.2/go.mod h1:Epx+iuz8sNs7mNKhxzH4fWXGNpZwUaJKRS1noLXviQ8= +github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= +github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= +github.com/ncruces/go-strftime v0.1.9 h1:bY0MQC28UADQmHmaF5dgpLmImcShSi2kHU9XLdhx/f4= +github.com/ncruces/go-strftime v0.1.9/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= +github.com/openai/openai-go/v2 v2.7.0 h1:/8MSFCXcasin7AyuWQ2au6FraXL71gzAs+VfbMv+J3k= +github.com/openai/openai-go/v2 v2.7.0/go.mod h1:jrJs23apqJKKbT+pqtFgNKpRju/KP9zpUTZhz3GElQE= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/prometheus/client_golang v1.23.2 h1:Je96obch5RDVy3FDMndoUsjAhG5Edi49h0RJWRi/o0o= +github.com/prometheus/client_golang v1.23.2/go.mod h1:Tb1a6LWHB3/SPIzCoaDXI4I8UHKeFTEQ1YCr+0Gyqmg= +github.com/prometheus/client_model v0.6.2 h1:oBsgwpGs7iVziMvrGhE53c/GrLUsZdHnqNwqPLxwZyk= +github.com/prometheus/client_model v0.6.2/go.mod h1:y3m2F6Gdpfy6Ut/GBsUqTWZqCUvMVzSfMLjcu6wAwpE= +github.com/prometheus/common v0.66.1 h1:h5E0h5/Y8niHc5DlaLlWLArTQI7tMrsfQjHV+d9ZoGs= +github.com/prometheus/common v0.66.1/go.mod h1:gcaUsgf3KfRSwHY4dIMXLPV0K/Wg1oZ8+SbZk/HH/dA= +github.com/prometheus/procfs v0.16.1 h1:hZ15bTNuirocR6u0JZ6BAHHmwS1p8B4P6MRqxtzMyRg= +github.com/prometheus/procfs v0.16.1/go.mod h1:teAbpZRB1iIAJYREa1LsoWUXykVXA1KlTmWl8x/U+Is= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= +github.com/rivo/uniseg v0.1.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= +github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= +github.com/rivo/uniseg v0.4.4 h1:8TfxU8dW6PdqD27gjM8MVNuicgxIjxpm4K7x4jp8sis= +github.com/rivo/uniseg v0.4.4/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= +github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII= +github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o= +github.com/spf13/cast v1.7.1 h1:cuNEagBQEHWN1FnbGEjCXL2szYEXqfJPbP2HNUaca9Y= +github.com/spf13/cast v1.7.1/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= +github.com/tidwall/gjson v1.18.0 h1:FIDeeyB800efLX89e5a8Y0BNH+LOngJyGrIWxG2FKQY= +github.com/tidwall/gjson v1.18.0/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= +github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= +github.com/tidwall/match v1.2.0 h1:0pt8FlkOwjN2fPt4bIl4BoNxb98gGHN2ObFEDkrfZnM= +github.com/tidwall/match v1.2.0/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= +github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= +github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4= +github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= +github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= +github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28= +github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= +github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= +github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= +github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= +go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA= +go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A= +go.opentelemetry.io/otel v1.33.0 h1:/FerN9bax5LoK51X/sI0SVYrjSE0/yUL7DpxW4K3FWw= +go.opentelemetry.io/otel v1.33.0/go.mod h1:SUUkR6csvUQl+yjReHu5uM3EtVV7MBm5FHKRlNx4I8I= +go.opentelemetry.io/otel/metric v1.33.0 h1:r+JOocAyeRVXD8lZpjdQjzMadVZp2M4WmQ+5WtEnklQ= +go.opentelemetry.io/otel/metric v1.33.0/go.mod h1:L9+Fyctbp6HFTddIxClbQkjtubW6O9QS3Ann/M82u6M= +go.opentelemetry.io/otel/sdk v1.16.0 h1:Z1Ok1YsijYL0CSJpHt4cS3wDDh7p572grzNrBMiMWgE= +go.opentelemetry.io/otel/sdk v1.16.0/go.mod h1:tMsIuKXuuIWPBAOrH+eHtvhTL+SntFtXF9QD68aP6p4= +go.opentelemetry.io/otel/trace v1.33.0 h1:cCJuF7LRjUFso9LPnEAHJDB2pqzp+hbO8eu1qqW2d/s= +go.opentelemetry.io/otel/trace v1.33.0/go.mod h1:uIcdVUZMpTAmz0tI1z04GoVSezK37CbGV4fr1f2nBck= +go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= +go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= +go.yaml.in/yaml/v2 v2.4.2 h1:DzmwEr2rDGHl7lsFgAHxmNz/1NlQ7xLIrlN2h5d1eGI= +go.yaml.in/yaml/v2 v2.4.2/go.mod h1:081UH+NErpNdqlCXm3TtEran0rJZGxAYx9hb/ELlsPU= +golang.org/x/exp v0.0.0-20250819193227-8b4c13bb791b h1:DXr+pvt3nC887026GRP39Ej11UATqWDmWuS99x26cD0= +golang.org/x/exp v0.0.0-20250819193227-8b4c13bb791b/go.mod h1:4QTo5u+SEIbbKW1RacMZq1YEfOBqeXa19JeshGi+zc4= +golang.org/x/mod v0.27.0 h1:kb+q2PyFnEADO2IEF935ehFUXlWiNjJWtRNgBLSfbxQ= +golang.org/x/mod v0.27.0/go.mod h1:rWI627Fq0DEoudcK+MBkNkCe0EetEaDSwJJkCcjpazc= +golang.org/x/net v0.43.0 h1:lat02VYK2j4aLzMzecihNvTlJNQUq316m2Mr9rnM6YE= +golang.org/x/net v0.43.0/go.mod h1:vhO1fvI4dGsIjh73sWfUVjj3N7CA9WkKJNQm2svM6Jg= +golang.org/x/sync v0.16.0 h1:ycBJEhp9p4vXvUZNszeOq0kGTPghopOL8q0fq3vstxw= +golang.org/x/sync v0.16.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA= +golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.36.0 h1:KVRy2GtZBrk1cBYA7MKu5bEZFxQk4NIDV6RLVcC8o0k= +golang.org/x/sys v0.36.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= +golang.org/x/term v0.34.0 h1:O/2T7POpk0ZZ7MAzMeWFSg6S5IpWd/RXDlM9hgM3DR4= +golang.org/x/term v0.34.0/go.mod h1:5jC53AEywhIVebHgPVeg0mj8OD3VO9OzclacVrqpaAw= +golang.org/x/text v0.28.0 h1:rhazDwis8INMIwQ4tpjLDzUhx6RlXqZNPEM0huQojng= +golang.org/x/text v0.28.0/go.mod h1:U8nCwOR8jO/marOQ0QbDiOngZVEBB7MAiitBuMjXiNU= +golang.org/x/tools v0.36.0 h1:kWS0uv/zsvHEle1LbV5LE8QujrxB3wfQyxHfhOk0Qkg= +golang.org/x/tools v0.36.0/go.mod h1:WBDiHKJK8YgLHlcQPYQzNCkUxUypCaa5ZegCVutKm+s= +golang.org/x/xerrors v0.0.0-20240716161551-93cc26a95ae9 h1:LLhsEBxRTBLuKlQxFBYUOU8xyFgXv6cOTp2HASDlsDk= +golang.org/x/xerrors v0.0.0-20240716161551-93cc26a95ae9/go.mod h1:NDW/Ps6MPRej6fsCIbMTohpP40sJ/P/vI1MoTEGwX90= +google.golang.org/genproto v0.0.0-20230726155614-23370e0ffb3e h1:xIXmWJ303kJCuogpj0bHq+dcjcZHU+XFyc1I0Yl9cRg= +google.golang.org/genproto v0.0.0-20230726155614-23370e0ffb3e/go.mod h1:0ggbjUrZYpy1q+ANUS30SEoGZ53cdfwtbuG7Ptgy108= +google.golang.org/genproto/googleapis/api v0.0.0-20240318140521-94a12d6c2237 h1:RFiFrvy37/mpSpdySBDrUdipW/dHwsRwh3J3+A9VgT4= +google.golang.org/genproto/googleapis/api v0.0.0-20240318140521-94a12d6c2237/go.mod h1:Z5Iiy3jtmioajWHDGFk7CeugTyHtPvMHA4UTmUkyalE= +google.golang.org/genproto/googleapis/rpc v0.0.0-20240722135656-d784300faade h1:oCRSWfwGXQsqlVdErcyTt4A93Y8fo0/9D4b1gnI++qo= +google.golang.org/genproto/googleapis/rpc v0.0.0-20240722135656-d784300faade/go.mod h1:Ue6ibwXGpU+dqIcODieyLOcgj7z8+IcskoNIgZxtrFY= +google.golang.org/grpc v1.64.1 h1:LKtvyfbX3UGVPFcGqJ9ItpVWW6oN/2XqTxfAnwRRXiA= +google.golang.org/grpc v1.64.1/go.mod h1:hiQF4LFZelK2WKaP6W0L92zGHtiQdZxk8CrSdvyjeP0= +google.golang.org/protobuf v1.36.8 h1:xHScyCOEuuwZEc6UtSOvPbAT4zRh0xcNRYekJwfqyMc= +google.golang.org/protobuf v1.36.8/go.mod h1:fuxRtAxBytpl4zzqUh6/eyUujkJdNiuEkXntxiD/uRU= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +modernc.org/cc/v4 v4.26.5 h1:xM3bX7Mve6G8K8b+T11ReenJOT+BmVqQj0FY5T4+5Y4= +modernc.org/cc/v4 v4.26.5/go.mod h1:uVtb5OGqUKpoLWhqwNQo/8LwvoiEBLvZXIQ/SmO6mL0= +modernc.org/ccgo/v4 v4.28.1 h1:wPKYn5EC/mYTqBO373jKjvX2n+3+aK7+sICCv4Fjy1A= +modernc.org/ccgo/v4 v4.28.1/go.mod h1:uD+4RnfrVgE6ec9NGguUNdhqzNIeeomeXf6CL0GTE5Q= +modernc.org/fileutil v1.3.40 h1:ZGMswMNc9JOCrcrakF1HrvmergNLAmxOPjizirpfqBA= +modernc.org/fileutil v1.3.40/go.mod h1:HxmghZSZVAz/LXcMNwZPA/DRrQZEVP9VX0V4LQGQFOc= +modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI= +modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito= +modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks= +modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI= +modernc.org/libc v1.66.10 h1:yZkb3YeLx4oynyR+iUsXsybsX4Ubx7MQlSYEw4yj59A= +modernc.org/libc v1.66.10/go.mod h1:8vGSEwvoUoltr4dlywvHqjtAqHBaw0j1jI7iFBTAr2I= +modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU= +modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg= +modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI= +modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw= +modernc.org/opt v0.1.4 h1:2kNGMRiUjrp4LcaPuLY2PzUfqM/w9N23quVwhKt5Qm8= +modernc.org/opt v0.1.4/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns= +modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w= +modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE= +modernc.org/sqlite v1.40.1 h1:VfuXcxcUWWKRBuP8+BR9L7VnmusMgBNNnBYGEe9w/iY= +modernc.org/sqlite v1.40.1/go.mod h1:9fjQZ0mB1LLP0GYrp39oOJXx/I2sxEnZtzCmEQIKvGE= +modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0= +modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A= +modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y= +modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM= diff --git a/example/main.go b/example/main.go new file mode 100644 index 0000000..7a9b4f7 --- /dev/null +++ b/example/main.go @@ -0,0 +1,115 @@ +// This is an example server demonstrating aibridge usage. +// Run with: go run *.go +package main + +import ( + "context" + "database/sql" + "net/http" + "os" + "regexp" + + "cdr.dev/slog" + "cdr.dev/slog/sloggers/sloghuman" + "github.com/coder/aibridge" + "github.com/coder/aibridge/mcp" + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promhttp" + + _ "modernc.org/sqlite" +) + +func main() { + ctx := context.Background() + logger := slog.Make(sloghuman.Sink(os.Stderr)).Leveled(slog.LevelDebug) + + // Initialize SQLite database with WAL mode for better concurrency. + db, err := sql.Open("sqlite", "aibridge.db?_journal_mode=WAL&_busy_timeout=5000") + if err != nil { + logger.Fatal(ctx, "open database", slog.Error(err)) + } + defer db.Close() + db.SetMaxOpenConns(1) // SQLite only supports one writer at a time. + + if err := initSchema(db); err != nil { + logger.Fatal(ctx, "init schema", slog.Error(err)) + } + + recorder, err := NewSQLiteRecorder(db, logger) + if err != nil { + logger.Fatal(ctx, "create recorder", slog.Error(err)) + } + defer recorder.Close() + + // Configure providers. + providers := []aibridge.Provider{ + aibridge.NewAnthropicProvider(aibridge.AnthropicConfig{ + Key: os.Getenv("ANTHROPIC_API_KEY"), + }, nil), + aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{ + Key: os.Getenv("OPENAI_API_KEY"), + }), + } + + // Setup metrics. + reg := prometheus.NewRegistry() + metrics := aibridge.NewMetrics(reg) + + // Optional: Configure MCP server for centralized tool injection. + // DeepWiki provides free access to public GitHub repo documentation. + // See: https://mcp.deepwiki.com + var mcpProxy mcp.ServerProxier + deepwikiProxy, err := mcp.NewStreamableHTTPServerProxy( + logger.Named("mcp.deepwiki"), + "deepwiki", // server name (tools prefixed as bmcp_deepwiki_*) + "https://mcp.deepwiki.com/mcp", // no auth required for public repos + nil, // headers + regexp.MustCompile(`^ask_question$`), // allowlist: only ask_question tool + nil, // denylist + ) + if err != nil { + logger.Fatal(ctx, "create deepwiki mcp proxy", slog.Error(err)) + } + + mcpProxy = mcp.NewServerProxyManager(map[string]mcp.ServerProxier{"deepwiki": deepwikiProxy}) + if err := mcpProxy.Init(ctx); err != nil { + logger.Warn(ctx, "mcp init warning", slog.Error(err)) + } + + // Create the bridge with SQLite recorder. + bridge, err := aibridge.NewRequestBridge( + ctx, + providers, + recorder, + mcpProxy, + metrics, + logger, + ) + if err != nil { + logger.Fatal(ctx, "create bridge", slog.Error(err)) + } + defer bridge.Shutdown(ctx) + + // Setup HTTP routes. + mux := http.NewServeMux() + mux.Handle("/metrics", promhttp.HandlerFor(reg, promhttp.HandlerOpts{})) + mux.Handle("/", actorMiddleware(bridge)) + + logger.Info(ctx, "listening on :8080") + if err := http.ListenAndServe(":8080", mux); err != nil { + logger.Fatal(ctx, "http server error", slog.Error(err)) + } +} + +// actorMiddleware injects actor identity into request context. +// In production, the user ID should be extracted from auth headers/tokens. +func actorMiddleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + userID := r.Header.Get("X-User-ID") + if userID == "" { + userID = "anonymous" + } + ctx := aibridge.AsActor(r.Context(), userID, nil) + next.ServeHTTP(w, r.WithContext(ctx)) + }) +} diff --git a/example/recorder.go b/example/recorder.go new file mode 100644 index 0000000..d2967eb --- /dev/null +++ b/example/recorder.go @@ -0,0 +1,147 @@ +package main + +import ( + "context" + "database/sql" + "encoding/json" + + "cdr.dev/slog" + "github.com/coder/aibridge" + "github.com/google/uuid" +) + +// SQLiteRecorder implements aibridge.Recorder and persists usage data to SQLite. +type SQLiteRecorder struct { + db *sql.DB + logger slog.Logger + + stmtInsertInterception *sql.Stmt + stmtUpdateInterception *sql.Stmt + stmtInsertTokenUsage *sql.Stmt + stmtInsertPromptUsage *sql.Stmt + stmtInsertToolUsage *sql.Stmt +} + +func NewSQLiteRecorder(db *sql.DB, logger slog.Logger) (*SQLiteRecorder, error) { + r := &SQLiteRecorder{db: db, logger: logger} + + var err error + r.stmtInsertInterception, err = db.Prepare(` + INSERT INTO aibridge_interceptions (id, initiator_id, provider, model, started_at, metadata) + VALUES (?, ?, ?, ?, ?, ?)`) + if err != nil { + return nil, err + } + + r.stmtUpdateInterception, err = db.Prepare(` + UPDATE aibridge_interceptions SET ended_at = ? WHERE id = ?`) + if err != nil { + return nil, err + } + + r.stmtInsertTokenUsage, err = db.Prepare(` + INSERT INTO aibridge_token_usages (id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?)`) + if err != nil { + return nil, err + } + + r.stmtInsertPromptUsage, err = db.Prepare(` + INSERT INTO aibridge_user_prompts (id, interception_id, provider_response_id, prompt, metadata, created_at) + VALUES (?, ?, ?, ?, ?, ?)`) + if err != nil { + return nil, err + } + + r.stmtInsertToolUsage, err = db.Prepare(` + INSERT INTO aibridge_tool_usages (id, interception_id, provider_response_id, server_url, tool, input, injected, invocation_error, metadata, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`) + if err != nil { + return nil, err + } + + return r, nil +} + +func (r *SQLiteRecorder) Close() error { + r.stmtInsertInterception.Close() + r.stmtUpdateInterception.Close() + r.stmtInsertTokenUsage.Close() + r.stmtInsertPromptUsage.Close() + r.stmtInsertToolUsage.Close() + return nil +} + +func (r *SQLiteRecorder) RecordInterception(ctx context.Context, req *aibridge.InterceptionRecord) error { + metadata, _ := json.Marshal(req.Metadata) + _, err := r.stmtInsertInterception.ExecContext(ctx, + req.ID, req.InitiatorID, req.Provider, req.Model, req.StartedAt, string(metadata), + ) + if err != nil { + r.logger.Warn(ctx, "failed to record interception", slog.Error(err)) + } + return err +} + +func (r *SQLiteRecorder) RecordInterceptionEnded(ctx context.Context, req *aibridge.InterceptionRecordEnded) error { + _, err := r.stmtUpdateInterception.ExecContext(ctx, req.EndedAt, req.ID) + if err != nil { + r.logger.Warn(ctx, "failed to record interception end", slog.Error(err)) + } + return err +} + +func (r *SQLiteRecorder) RecordTokenUsage(ctx context.Context, req *aibridge.TokenUsageRecord) error { + // Build metadata, merging extra token types. + merged := make(map[string]any) + for k, v := range req.Metadata { + merged[k] = v + } + for k, v := range req.ExtraTokenTypes { + merged[k] = v + } + metadata, _ := json.Marshal(merged) + + _, err := r.stmtInsertTokenUsage.ExecContext(ctx, + uuid.NewString(), req.InterceptionID, req.MsgID, req.Input, req.Output, string(metadata), req.CreatedAt, + ) + if err != nil { + r.logger.Warn(ctx, "failed to record token usage", slog.Error(err)) + } + return err +} + +func (r *SQLiteRecorder) RecordPromptUsage(ctx context.Context, req *aibridge.PromptUsageRecord) error { + metadata, _ := json.Marshal(req.Metadata) + _, err := r.stmtInsertPromptUsage.ExecContext(ctx, + uuid.NewString(), req.InterceptionID, req.MsgID, req.Prompt, string(metadata), req.CreatedAt, + ) + if err != nil { + r.logger.Warn(ctx, "failed to record prompt usage", slog.Error(err)) + } + return err +} + +func (r *SQLiteRecorder) RecordToolUsage(ctx context.Context, req *aibridge.ToolUsageRecord) error { + metadata, _ := json.Marshal(req.Metadata) + input, _ := json.Marshal(req.Args) + + var serverURL *string + if req.ServerURL != nil { + serverURL = req.ServerURL + } + + var invocationError *string + if req.InvocationError != nil { + errStr := req.InvocationError.Error() + invocationError = &errStr + } + + _, err := r.stmtInsertToolUsage.ExecContext(ctx, + uuid.NewString(), req.InterceptionID, req.MsgID, serverURL, req.Tool, string(input), req.Injected, invocationError, string(metadata), req.CreatedAt, + ) + if err != nil { + r.logger.Warn(ctx, "failed to record tool usage", slog.Error(err)) + } + return err +} diff --git a/example/schema.go b/example/schema.go new file mode 100644 index 0000000..790de16 --- /dev/null +++ b/example/schema.go @@ -0,0 +1,61 @@ +package main + +import "database/sql" + +// initSchema creates the aibridge tables following the coder/coder data model. +// See: https://github.com/coder/coder/blob/c87c33f7dda82eb91ee8ba9504f749101bb367d6/coderd/database/dump.sql#L1052-L1115 +func initSchema(db *sql.DB) error { + schema := ` + CREATE TABLE IF NOT EXISTS aibridge_interceptions ( + id TEXT PRIMARY KEY, + initiator_id TEXT NOT NULL, + provider TEXT NOT NULL, + model TEXT NOT NULL, + started_at DATETIME NOT NULL, + ended_at DATETIME, + metadata TEXT + ); + + CREATE TABLE IF NOT EXISTS aibridge_token_usages ( + id TEXT PRIMARY KEY, + interception_id TEXT NOT NULL, + provider_response_id TEXT NOT NULL, + input_tokens INTEGER NOT NULL, + output_tokens INTEGER NOT NULL, + metadata TEXT, + created_at DATETIME NOT NULL, + FOREIGN KEY (interception_id) REFERENCES aibridge_interceptions(id) + ); + + CREATE TABLE IF NOT EXISTS aibridge_user_prompts ( + id TEXT PRIMARY KEY, + interception_id TEXT NOT NULL, + provider_response_id TEXT NOT NULL, + prompt TEXT NOT NULL, + metadata TEXT, + created_at DATETIME NOT NULL, + FOREIGN KEY (interception_id) REFERENCES aibridge_interceptions(id) + ); + + CREATE TABLE IF NOT EXISTS aibridge_tool_usages ( + id TEXT PRIMARY KEY, + interception_id TEXT NOT NULL, + provider_response_id TEXT NOT NULL, + server_url TEXT, + tool TEXT NOT NULL, + input TEXT NOT NULL, + injected BOOLEAN NOT NULL DEFAULT FALSE, + invocation_error TEXT, + metadata TEXT, + created_at DATETIME NOT NULL, + FOREIGN KEY (interception_id) REFERENCES aibridge_interceptions(id) + ); + + CREATE INDEX IF NOT EXISTS idx_interceptions_initiator ON aibridge_interceptions(initiator_id); + CREATE INDEX IF NOT EXISTS idx_token_usages_interception ON aibridge_token_usages(interception_id); + CREATE INDEX IF NOT EXISTS idx_user_prompts_interception ON aibridge_user_prompts(interception_id); + CREATE INDEX IF NOT EXISTS idx_tool_usages_interception ON aibridge_tool_usages(interception_id); + ` + _, err := db.Exec(schema) + return err +}