diff --git a/.configs/sqlc.yaml b/.configs/sqlc.yaml index e43f0c06f..da1ec923b 100644 --- a/.configs/sqlc.yaml +++ b/.configs/sqlc.yaml @@ -219,3 +219,12 @@ sql: <<: *default_go package: "restteamsapisql" out: "../internal/rest/restteamsapi/restteamsapisql" + + - <<: *default_domain + name: "Agent SQL" + queries: "../internal/agent/queries" + gen: + go: + <<: *default_go + package: "agentsql" + out: "../internal/agent/agentsql" diff --git a/.env.example b/.env.example index 3b639e6fa..933c52a0f 100644 --- a/.env.example +++ b/.env.example @@ -44,3 +44,11 @@ AIVEN_PROJECTS='{"dev": {"id":"nav-dev","vpc":"somevpc","endpoint_id":"some_endp #Uncomment if you want to use the github.com/nais/v13s api locally #VULNERABILITIES_ENDPOINT=localhost:50051 #VULNERABILITIES_SERVICE_ACCOUNT=notused + +AGENT_ENABLED=false +AGENT_RAG_DUCKDB_PATH=./data/rag_index.duckdb +AGENT_VERTEX_AI_PROJECT_ID=naissistant-test +AGENT_VERTEX_AI_LOCATION=europe-north1 +AGENT_VERTEX_AI_MODEL=gemini-2.5-pro +AGENT_VERTEX_AI_EMBEDDING_MODEL=gemini-embedding-001 +AGENT_VERTEX_AI_INCLUDE_THOUGHTS=true diff --git a/.gitignore b/.gitignore index aecfdb698..2623e86a4 100644 --- a/.gitignore +++ b/.gitignore @@ -14,3 +14,4 @@ hack/coverprofile.txt result* data/k8s/dev/tbd/ data/k8s/prod-gcp/ +data/*.duckdb diff --git a/.mise-tasks/build/gen-rag-index b/.mise-tasks/build/gen-rag-index new file mode 100755 index 000000000..d030c19ab --- /dev/null +++ b/.mise-tasks/build/gen-rag-index @@ -0,0 +1,6 @@ +#!/usr/bin/env sh +#MISE description="Build gen_rag_index" + +set -e + +go build -o bin/gen_rag_index ./cmd/gen_rag_index diff --git a/.mise-tasks/generate/rag-index b/.mise-tasks/generate/rag-index new file mode 100755 index 000000000..01765ca7c --- /dev/null +++ b/.mise-tasks/generate/rag-index @@ -0,0 +1,7 @@ +#!/bin/bash +#MISE description="Generate DuckDB RAG index from documentation search index" +#MISE depends=["build:gen-rag-index"] + +set -e + +go run ./cmd/gen_rag_index diff --git a/cmd/gen_rag_index/main.go b/cmd/gen_rag_index/main.go new file mode 100644 index 000000000..9730f12a4 --- /dev/null +++ b/cmd/gen_rag_index/main.go @@ -0,0 +1,12 @@ +package main + +import ( + "context" + + genragindex "github.com/nais/api/internal/cmd/gen_rag_index" +) + +func main() { + ctx := context.Background() + genragindex.Run(ctx) +} diff --git a/go.mod b/go.mod index d9feb8aad..25bdf84e0 100644 --- a/go.mod +++ b/go.mod @@ -27,6 +27,7 @@ require ( github.com/bombsimon/logrusr/v4 v4.1.0 github.com/btcsuite/btcutil v1.0.2 github.com/coreos/go-oidc/v3 v3.12.0 + github.com/duckdb/duckdb-go/v2 v2.5.4 github.com/evanphx/json-patch/v5 v5.9.11 github.com/exaring/otelpgx v0.9.0 github.com/go-chi/chi/v5 v5.2.4 @@ -73,12 +74,13 @@ require ( go.opentelemetry.io/otel/sdk v1.41.0 go.opentelemetry.io/otel/sdk/metric v1.41.0 go.opentelemetry.io/otel/trace v1.41.0 - golang.org/x/exp v0.0.0-20250808145144-a408d31f581a + golang.org/x/exp v0.0.0-20251209150349-8475f28825e9 golang.org/x/oauth2 v0.36.0 golang.org/x/sync v0.19.0 golang.org/x/text v0.34.0 golang.org/x/tools v0.42.0 google.golang.org/api v0.269.0 + google.golang.org/genai v1.42.0 google.golang.org/genproto/googleapis/api v0.0.0-20260128011058-8636f8732409 google.golang.org/grpc v1.79.3 google.golang.org/protobuf v1.36.11 @@ -118,6 +120,7 @@ require ( github.com/agnivade/levenshtein v1.2.1 // indirect github.com/alecthomas/units v0.0.0-20240927000941-0f3dac36c52b // indirect github.com/antlr4-go/antlr/v4 v4.13.1 // indirect + github.com/apache/arrow-go/v18 v18.4.1 // indirect github.com/apache/arrow/go/v15 v15.0.2 // indirect github.com/apapsch/go-jsonmerge/v2 v2.0.0 // indirect github.com/armon/go-metrics v0.4.1 // indirect @@ -179,6 +182,14 @@ require ( github.com/docker/docker v28.5.2+incompatible // indirect github.com/docker/go-connections v0.5.0 // indirect github.com/docker/go-units v0.5.0 // indirect + github.com/duckdb/duckdb-go-bindings v0.1.24 // indirect + github.com/duckdb/duckdb-go-bindings/darwin-amd64 v0.1.24 // indirect + github.com/duckdb/duckdb-go-bindings/darwin-arm64 v0.1.24 // indirect + github.com/duckdb/duckdb-go-bindings/linux-amd64 v0.1.24 // indirect + github.com/duckdb/duckdb-go-bindings/linux-arm64 v0.1.24 // indirect + github.com/duckdb/duckdb-go-bindings/windows-amd64 v0.1.24 // indirect + github.com/duckdb/duckdb-go/arrowmapping v0.0.27 // indirect + github.com/duckdb/duckdb-go/mapping v0.0.27 // indirect github.com/dustin/go-humanize v1.0.1 // indirect github.com/ebitengine/purego v0.9.0 // indirect github.com/edsrzf/mmap-go v1.2.0 // indirect @@ -234,7 +245,7 @@ require ( github.com/golang/snappy v1.0.0 // indirect github.com/google/btree v1.1.3 // indirect github.com/google/cel-go v0.26.0 // indirect - github.com/google/flatbuffers v25.2.10+incompatible // indirect + github.com/google/flatbuffers v25.9.23+incompatible // indirect github.com/google/generative-ai-go v0.19.0 // indirect github.com/google/gnostic-models v0.7.1 // indirect github.com/google/s2a-go v0.1.9 // indirect @@ -283,7 +294,7 @@ require ( github.com/julienschmidt/httprouter v1.3.0 // indirect github.com/kamstrup/intmap v0.5.1 // indirect github.com/kelseyhightower/envconfig v1.4.0 // indirect - github.com/klauspost/compress v1.18.1 // indirect + github.com/klauspost/compress v1.18.2 // indirect github.com/klauspost/cpuid/v2 v2.3.0 // indirect github.com/knadh/koanf/maps v0.1.2 // indirect github.com/knadh/koanf/parsers/yaml v0.1.0 // indirect @@ -376,7 +387,7 @@ require ( github.com/spf13/cobra v1.10.2 // indirect github.com/spf13/pflag v1.0.10 // indirect github.com/sqlc-dev/sqlc v1.29.0 // indirect - github.com/stoewer/go-strcase v1.3.0 // indirect + github.com/stoewer/go-strcase v1.3.1 // indirect github.com/stretchr/objx v0.5.3 // indirect github.com/tetratelabs/wazero v1.9.0 // indirect github.com/tjhop/slog-gokit v0.1.4 // indirect diff --git a/go.sum b/go.sum index df2100fb7..61f8325ff 100644 --- a/go.sum +++ b/go.sum @@ -112,12 +112,18 @@ github.com/alicebob/miniredis/v2 v2.35.0 h1:QwLphYqCEAo1eu1TqPRN2jgVMPBweeQcR21j github.com/alicebob/miniredis/v2 v2.35.0/go.mod h1:TcL7YfarKPGDAthEtl5NBeHZfeUQj6OXMm/+iu5cLMM= github.com/andreyvit/diff v0.0.0-20170406064948-c7f18ee00883 h1:bvNMNQO63//z+xNgfBlViaCIJKLlCJ6/fmUseuG0wVQ= github.com/andreyvit/diff v0.0.0-20170406064948-c7f18ee00883/go.mod h1:rCTlJbsFo29Kk6CurOXKm700vrz8f0KW0JNfpkRJY/8= +github.com/andybalholm/brotli v1.2.0 h1:ukwgCxwYrmACq68yiUqwIWnGY0cTPox/M94sVwToPjQ= +github.com/andybalholm/brotli v1.2.0/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY= github.com/andybalholm/cascadia v1.3.3 h1:AG2YHrzJIm4BZ19iwJ/DAua6Btl3IwJX+VI4kktS1LM= github.com/andybalholm/cascadia v1.3.3/go.mod h1:xNd9bqTn98Ln4DwST8/nG+H0yuB8Hmgu1YHNnWw0GeA= github.com/antlr4-go/antlr/v4 v4.13.1 h1:SqQKkuVZ+zWkMMNkjy5FZe5mr5WURWnlpmOuzYWrPrQ= github.com/antlr4-go/antlr/v4 v4.13.1/go.mod h1:GKmUxMtwp6ZgGwZSva4eWPC5mS6vUAmOABFgjdkM7Nw= +github.com/apache/arrow-go/v18 v18.4.1 h1:q/jVkBWCJOB9reDgaIZIdruLQUb1kbkvOnOFezVH1C4= +github.com/apache/arrow-go/v18 v18.4.1/go.mod h1:tLyFubsAl17bvFdUAy24bsSvA/6ww95Iqi67fTpGu3E= github.com/apache/arrow/go/v15 v15.0.2 h1:60IliRbiyTWCWjERBCkO1W4Qun9svcYoZrSLcyOsMLE= github.com/apache/arrow/go/v15 v15.0.2/go.mod h1:DGXsR3ajT524njufqf95822i+KTh+yea1jass9YXgjA= +github.com/apache/thrift v0.22.0 h1:r7mTJdj51TMDe6RtcmNdQxgn9XcyfGDOzegMDRg47uc= +github.com/apache/thrift v0.22.0/go.mod h1:1e7J/O1Ae6ZQMTYdy9xa3w9k+XHWPfRvdPyJeynQ+/g= github.com/apapsch/go-jsonmerge/v2 v2.0.0 h1:axGnT1gRIfimI7gJifB699GoE/oq+F2MU7Dml6nw9rQ= github.com/apapsch/go-jsonmerge/v2 v2.0.0/go.mod h1:lvDnEdqiQrp0O42VQGgmlKpxL1AP2+08jFMw88y4klk= github.com/arbovm/levenshtein v0.0.0-20160628152529-48b4e1c0c4d0 h1:jfIu9sQUG6Ig+0+Ap1h4unLjW6YQJpKZVmUzxsD4E/Q= @@ -301,6 +307,24 @@ github.com/docker/go-connections v0.5.0 h1:USnMq7hx7gwdVZq1L49hLXaFtUdTADjXGp+uj github.com/docker/go-connections v0.5.0/go.mod h1:ov60Kzw0kKElRwhNs9UlUHAE/F9Fe6GLaXnqyDdmEXc= github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4= github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk= +github.com/duckdb/duckdb-go-bindings v0.1.24 h1:p1v3GruGHGcZD69cWauH6QrOX32oooqdUAxrWK3Fo6o= +github.com/duckdb/duckdb-go-bindings v0.1.24/go.mod h1:WA7U/o+b37MK2kiOPPueVZ+FIxt5AZFCjszi8hHeH18= +github.com/duckdb/duckdb-go-bindings/darwin-amd64 v0.1.24 h1:XhqMj+bvpTIm+hMeps1Kk94r2eclAswk2ISFs4jMm+g= +github.com/duckdb/duckdb-go-bindings/darwin-amd64 v0.1.24/go.mod h1:jfbOHwGZqNCpMAxV4g4g5jmWr0gKdMvh2fGusPubxC4= +github.com/duckdb/duckdb-go-bindings/darwin-arm64 v0.1.24 h1:OyHr5PykY5FG81jchpRoESMDQX1HK66PdNsfxoHxbwM= +github.com/duckdb/duckdb-go-bindings/darwin-arm64 v0.1.24/go.mod h1:zLVtv1a7TBuTPvuAi32AIbnuw7jjaX5JElZ+urv1ydc= +github.com/duckdb/duckdb-go-bindings/linux-amd64 v0.1.24 h1:6Y4VarmcT7Oe8stwta4dOLlUX8aG4ciG9VhFKnp91a4= +github.com/duckdb/duckdb-go-bindings/linux-amd64 v0.1.24/go.mod h1:GCaBoYnuLZEva7BXzdXehTbqh9VSvpLB80xcmxGBGs8= +github.com/duckdb/duckdb-go-bindings/linux-arm64 v0.1.24 h1:NCAGH7o1RsJv631EQGOqs94ABtmYZO6JjMHkv7GIgG8= +github.com/duckdb/duckdb-go-bindings/linux-arm64 v0.1.24/go.mod h1:kpQSpJmDSSZQ3ikbZR1/8UqecqMeUkWFjFX2xZxlCuI= +github.com/duckdb/duckdb-go-bindings/windows-amd64 v0.1.24 h1:JOupXaHMMu8zLgq7v9uxPjl1CXSJHlISCxopMiqtkzU= +github.com/duckdb/duckdb-go-bindings/windows-amd64 v0.1.24/go.mod h1:wa+egSGXTPS16NPADFCK1yFyt3VSXxUS6Pt2fLnvRPM= +github.com/duckdb/duckdb-go/arrowmapping v0.0.27 h1:w0XKX+EJpAN4XOQlKxSxSKZq/tCVbRfTRBp98jA0q8M= +github.com/duckdb/duckdb-go/arrowmapping v0.0.27/go.mod h1:VkFx49Icor1bbxOPxAU8jRzwL0nTXICOthxVq4KqOqQ= +github.com/duckdb/duckdb-go/mapping v0.0.27 h1:QEta+qPEKmfhd89U8vnm4MVslj1UscmkyJwu8x+OtME= +github.com/duckdb/duckdb-go/mapping v0.0.27/go.mod h1:7C4QWJWG6UOV9b0iWanfF5ML1ivJPX45Kz+VmlvRlTA= +github.com/duckdb/duckdb-go/v2 v2.5.4 h1:+ip+wPCwf7Eu/dXxp19aLCxwpLUaeOy2UV/peBphXK0= +github.com/duckdb/duckdb-go/v2 v2.5.4/go.mod h1:CeobOFmWpf7MTDb+MW08/zIWP8TQ2jbPbMgGo5761tY= 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/ebitengine/purego v0.9.0 h1:mh0zpKBIXDceC63hpvPuGLiJ8ZAa3DfrFTudmfi8A4k= @@ -479,8 +503,8 @@ github.com/google/btree v1.1.3 h1:CVpQJjYgC4VbzxeGVHfvZrv1ctoYCAI8vbl07Fcxlyg= github.com/google/btree v1.1.3/go.mod h1:qOPhT0dTNdNzV6Z/lhRX0YXUafgPLFUh+gZMl761Gm4= github.com/google/cel-go v0.26.0 h1:DPGjXackMpJWH680oGY4lZhYjIameYmR+/6RBdDGmaI= github.com/google/cel-go v0.26.0/go.mod h1:A9O8OU9rdvrK5MQyrqfIxo1a0u4g3sF8KB6PUIaryMM= -github.com/google/flatbuffers v25.2.10+incompatible h1:F3vclr7C3HpB1k9mxCGRMXq6FdUalZ6H/pNX4FP1v0Q= -github.com/google/flatbuffers v25.2.10+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8= +github.com/google/flatbuffers v25.9.23+incompatible h1:rGZKv+wOb6QPzIdkM2KxhBZCDrA0DeN6DNmRDrqIsQU= +github.com/google/flatbuffers v25.9.23+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8= github.com/google/generative-ai-go v0.19.0 h1:R71szggh8wHMCUlEMsW2A/3T+5LdEIkiaHSYgSpUgdg= github.com/google/generative-ai-go v0.19.0/go.mod h1:JYolL13VG7j79kM5BtHz4qwONHkeJQzOCkKXnpqtS/E= github.com/google/gnostic-models v0.7.1 h1:SisTfuFKJSKM5CPZkffwi6coztzzeYUhc3v4yxLWH8c= @@ -667,8 +691,10 @@ github.com/keybase/go-keychain v0.0.1/go.mod h1:PdEILRW3i9D8JcdM+FmY6RwkHGnhHxXw github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8= github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck= github.com/kkdai/bstream v0.0.0-20161212061736-f391b8402d23/go.mod h1:J+Gs4SYgM6CZQHDETBtE9HaSEkGmuNXF86RwHhHUvq4= -github.com/klauspost/compress v1.18.1 h1:bcSGx7UbpBqMChDtsF28Lw6v/G94LPrrbMbdC3JH2co= -github.com/klauspost/compress v1.18.1/go.mod h1:ZQFFVG+MdnR0P+l6wpXgIL4NTtwiKIdBnrBd8Nrxr+0= +github.com/klauspost/asmfmt v1.3.2 h1:4Ri7ox3EwapiOjCki+hw14RyKk201CN4rzyCJRFLpK4= +github.com/klauspost/asmfmt v1.3.2/go.mod h1:AG8TuvYojzulgDAMCnYn50l/5QV3Bs/tp6j0HLHbNSE= +github.com/klauspost/compress v1.18.2 h1:iiPHWW0YrcFgpBYhsA6D1+fqHssJscY/Tm/y2Uqnapk= +github.com/klauspost/compress v1.18.2/go.mod h1:R0h/fSBs8DE4ENlcrlib3PsXS61voFxhIs2DeRhCvJ4= github.com/klauspost/cpuid/v2 v2.3.0 h1:S4CRMLnYUhGeDFDqkGriYKdfoFlDnMtqTiI/sFzhA9Y= github.com/klauspost/cpuid/v2 v2.3.0/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= github.com/knadh/koanf/maps v0.1.2 h1:RBfmAW5CnZT+PJ1CVc1QSJKf4Xu9kxfQgYVQSu8hpbo= @@ -745,6 +771,10 @@ github.com/microsoft/go-mssqldb v1.9.2 h1:nY8TmFMQOHpm2qVWo6y4I2mAmVdZqlGiMGAYt6 github.com/microsoft/go-mssqldb v1.9.2/go.mod h1:GBbW9ASTiDC+mpgWDGKdm3FnFLTUsLYN3iFL90lQ+PA= github.com/miekg/dns v1.1.68 h1:jsSRkNozw7G/mnmXULynzMNIsgY2dHC8LO6U6Ij2JEA= github.com/miekg/dns v1.1.68/go.mod h1:fujopn7TB3Pu3JM69XaawiU0wqjpL9/8xGop5UrTPps= +github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 h1:AMFGa4R4MiIpspGNG7Z948v4n35fFGB3RR3G/ry4FWs= +github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8/go.mod h1:mC1jAcsrzbxHt8iiaC+zU4b1ylILSosueou12R++wfY= +github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3 h1:+n/aFZefKZp7spd8DFdX7uMikMLXX4oubIzJF4kv/wI= +github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3/go.mod h1:RagcQ7I8IeTMnF8JTXieKnO4Z6JCsikNEzj0DwauVzE= github.com/mitchellh/copystructure v1.2.0 h1:vpKXTN4ewci03Vljg/q9QvCGUDttBOGBIa15WveJJGw= github.com/mitchellh/copystructure v1.2.0/go.mod h1:qLl+cE2AmVv+CoeAwDPye/v+N2HKCj9FbZEVFJRxO9s= github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y= @@ -990,8 +1020,8 @@ github.com/sqlc-dev/sqlc v1.29.0 h1:HQctoD7y/i29Bao53qXO7CZ/BV9NcvpGpsJWvz9nKWs= github.com/sqlc-dev/sqlc v1.29.0/go.mod h1:BavmYw11px5AdPOjAVHmb9fctP5A8GTziC38wBF9tp0= github.com/stackitcloud/stackit-sdk-go/core v0.17.3 h1:GsZGmRRc/3GJLmCUnsZswirr5wfLRrwavbnL/renOqg= github.com/stackitcloud/stackit-sdk-go/core v0.17.3/go.mod h1:HBCXJGPgdRulplDzhrmwC+Dak9B/x0nzNtmOpu+1Ahg= -github.com/stoewer/go-strcase v1.3.0 h1:g0eASXYtp+yvN9fK8sH94oCIk0fau9uV1/ZdJ0AVEzs= -github.com/stoewer/go-strcase v1.3.0/go.mod h1:fAH5hQ5pehh+j3nZfvwdk2RgEgQjAoM8wodgtPmh1xo= +github.com/stoewer/go-strcase v1.3.1 h1:iS0MdW+kVTxgMoE1LAZyMiYJFKlOzLooE4MxjirtkAs= +github.com/stoewer/go-strcase v1.3.1/go.mod h1:fAH5hQ5pehh+j3nZfvwdk2RgEgQjAoM8wodgtPmh1xo= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= @@ -1231,8 +1261,8 @@ golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL golang.org/x/exp v0.0.0-20190125153040-c74c464bbbf2/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= golang.org/x/exp v0.0.0-20190306152737-a1d7652674e8/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= golang.org/x/exp v0.0.0-20191030013958-a1ab85dbe136/go.mod h1:JXzH8nQsPlswgeRAPE3MuO9GYsAcnJvJ4vnMwN/5qkY= -golang.org/x/exp v0.0.0-20250808145144-a408d31f581a h1:Y+7uR/b1Mw2iSXZ3G//1haIiSElDQZ8KWh0h+sZPG90= -golang.org/x/exp v0.0.0-20250808145144-a408d31f581a/go.mod h1:rT6SFzZ7oxADUDx58pcaKFTcZ+inxAa9fTrYx/uVYwg= +golang.org/x/exp v0.0.0-20251209150349-8475f28825e9 h1:MDfG8Cvcqlt9XXrmEiD4epKn7VJHZO84hejP9Jmp0MM= +golang.org/x/exp v0.0.0-20251209150349-8475f28825e9/go.mod h1:EPRbTFwzwjXj9NpYyyrvenVh9Y+GFeEvMNh7Xuz7xgU= golang.org/x/exp/typeparams v0.0.0-20251209150349-8475f28825e9 h1:DXiKAjbw2KpfWz1Bq2YqF/dBDPEZGJsl3IA2JuVzy8U= golang.org/x/exp/typeparams v0.0.0-20251209150349-8475f28825e9/go.mod h1:4Mzdyp/6jzw9auFDJ3OMF5qksa7UvPnzKqTVGcb04ms= golang.org/x/image v0.0.0-20180708004352-c73c2afc3b81/go.mod h1:ux5Hcp/YLpHSI86hEcLt0YII63i6oz57MZXIpbrjZUs= @@ -1362,6 +1392,8 @@ google.golang.org/api v0.269.0 h1:qDrTOxKUQ/P0MveH6a7vZ+DNHxJQjtGm/uvdbdGXCQg= google.golang.org/api v0.269.0/go.mod h1:N8Wpcu23Tlccl0zSHEkcAZQKDLdquxK+l9r2LkwAauE= google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM= google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4= +google.golang.org/genai v1.42.0 h1:XFHfo0DDCzdzQALZoFs6nowAHO2cE95XyVvFLNaFLRY= +google.golang.org/genai v1.42.0/go.mod h1:A3kkl0nyBjyFlNjgxIwKq70julKbIxpSxqKO5gw/gmk= google.golang.org/genproto v0.0.0-20180518175338-11a468237815/go.mod h1:JiN7NxoALGmiZfu7CAH4rXhgtRTLTxftemlI0sWmxmc= google.golang.org/genproto v0.0.0-20180817151627-c66870c02cf8/go.mod h1:JiN7NxoALGmiZfu7CAH4rXhgtRTLTxftemlI0sWmxmc= google.golang.org/genproto v0.0.0-20190819201941-24fa4b261c55/go.mod h1:DMBHOl98Agz4BDEuKkezgsaosCRResVns1a3J2ZsMNc= diff --git a/internal/agent/CLIENT.md b/internal/agent/CLIENT.md new file mode 100644 index 000000000..7027e299f --- /dev/null +++ b/internal/agent/CLIENT.md @@ -0,0 +1,548 @@ +# Agent Chat API - Client Documentation + +This document describes the REST API for the Nais Agent chat service. + +## Base Path + +All agent endpoints are available under `/agent`. + +## Authentication + +The agent API uses the same authentication as the GraphQL API: +- Session cookies (OAuth2) +- JWT tokens via `Authorization: Bearer ` header +- API keys via `Authorization: Bearer ` header + +All endpoints require authentication. + +--- + +## Endpoints + +### POST /agent/chat + +Non-streaming chat endpoint. Sends a message and waits for the complete response. + +**Request Body:** +```json +{ + "message": "How do I deploy my application?", + "conversation_id": "optional-uuid-to-continue-conversation", + "context": { + "path": "/team/my-team/app/my-app", + "team": "my-team", + "app": "my-app", + "env": "dev" + } +} +``` + +| Field | Type | Required | Description | +|-------|------|----------|-------------| +| `message` | string | Yes | The user's message | +| `conversation_id` | string | No | UUID of an existing conversation to continue | +| `context` | object | No | Current UI context to help the assistant | +| `context.path` | string | No | Current page path | +| `context.team` | string | No | Current team slug | +| `context.app` | string | No | Current application name | +| `context.env` | string | No | Current environment | + +**Response:** +```json +{ + "conversation_id": "550e8400-e29b-41d4-a716-446655440000", + "message_id": "6ba7b810-9dad-11d1-80b4-00c04fd430c8", + "content": "To deploy your application, you need to...", + "blocks": [ + { + "type": "text", + "text": "To deploy your application, you need to..." + } + ], + "sources": [ + { + "title": "Deploying Applications", + "url": "https://docs.nais.io/deploy/" + } + ], + "usage": { + "input_tokens": 150, + "output_tokens": 75, + "total_tokens": 225 + } +} +``` + +| Field | Type | Description | +|-------|------|-------------| +| `conversation_id` | string | UUID of the conversation (new or existing) | +| `message_id` | string | UUID for this specific response message | +| `content` | string | Plain text of the response (concatenation of all text blocks) | +| `blocks` | array | Ordered content blocks — see [Content Blocks](#content-blocks) | +| `sources` | array | Documentation sources retrieved for this query | +| `sources[].title` | string | Document title | +| `sources[].url` | string | Document URL | +| `usage.input_tokens` | int | Tokens in the request | +| `usage.output_tokens` | int | Tokens in the response | +| `usage.total_tokens` | int | Total tokens used | + +--- + +### POST /agent/chat/stream + +Streaming chat endpoint using Server-Sent Events (SSE). Sends a message and streams the response in real-time. + +**Request Body:** Same as `POST /agent/chat`. + +**Response:** `Content-Type: text/event-stream` + +Each line is a Server-Sent Event in the format `data: \n\n`. Parse each `data:` value as JSON and handle it by `type`. + +#### Event Types + +**`metadata`** — Sent first. Contains the conversation and message IDs for this turn. +```json +{ + "type": "metadata", + "conversation_id": "550e8400-e29b-41d4-a716-446655440000", + "message_id": "6ba7b810-9dad-11d1-80b4-00c04fd430c8" +} +``` + +**`content`** — A chunk of the assistant's text response. Concatenate these in order to build the full response text. +```json +{ + "type": "content", + "content": "To deploy" +} +``` + +**`thinking`** — The model's reasoning/thought process (only when thinking mode is enabled on the server). See [Thinking Mode](#thinking-mode). +```json +{ + "type": "thinking", + "thinking": "The user is asking about deployment. I should look up their applications first." +} +``` + +**`tool_start`** — A tool call has begun. Use this to show a "working…" indicator in the UI. +```json +{ + "type": "tool_start", + "tool_call_id": "call_123456", + "tool_name": "execute_graphql", + "description": "Executing execute_graphql..." +} +``` + +**`tool_end`** — A tool call has finished. +```json +{ + "type": "tool_end", + "tool_call_id": "call_123456", + "tool_name": "execute_graphql", + "description": "Executed execute_graphql", + "tool_success": true +} +``` + +**`chart`** — A Prometheus chart to be rendered by the client. See [Rendering Charts](#rendering-charts). +```json +{ + "type": "chart", + "chart": { + "chart_type": "line", + "title": "CPU Usage for my-app", + "environment": "dev", + "query": "sum(rate(container_cpu_usage_seconds_total{app=\"my-app\"}[5m])) by (pod)", + "interval": "1h", + "y_format": "cpu_cores", + "label_template": "{pod}" + } +} +``` + +**`sources`** — Documentation sources retrieved for this query. Sent once near the end of the stream, before `done`. Only present if the RAG search returned results. +```json +{ + "type": "sources", + "sources": [ + { + "title": "Deploying Applications", + "url": "https://docs.nais.io/deploy/" + } + ] +} +``` + +**`usage`** — Token usage for the full turn. Sent once near the end of the stream, before `done`. +```json +{ + "type": "usage", + "usage": { + "input_tokens": 150, + "output_tokens": 75, + "total_tokens": 225, + "max_tokens": 131072 + } +} +``` + +**`done`** — The stream is complete. Stop processing after receiving this. +```json +{ + "type": "done" +} +``` + +**`error`** — A fatal error occurred. Stop processing after receiving this. +```json +{ + "type": "error", + "error_code": "stream_error", + "error_message": "Description of the error" +} +``` + +#### Event Order + +A complete stream always follows this order: + +1. `metadata` (always first) +2. Zero or more cycles of: `thinking`?, `tool_start`, `tool_end`, `chart`? +3. `thinking`? chunks (if the model reasons before replying) +4. `content` chunks (the actual response text) +5. `usage` (once, before done) +6. `sources`? (once, only if RAG returned results) +7. `done` + +If an `error` event is received at any point, the stream ends immediately. + +#### Example: Simple response + +``` +data: {"type":"metadata","conversation_id":"550e8400-e29b-41d4-a716-446655440000","message_id":"6ba7b810-9dad-11d1-80b4-00c04fd430c8"} + +data: {"type":"tool_start","tool_call_id":"call_123","tool_name":"execute_graphql","description":"Executing execute_graphql..."} + +data: {"type":"tool_end","tool_call_id":"call_123","tool_name":"execute_graphql","description":"Executed execute_graphql","tool_success":true} + +data: {"type":"content","content":"Based on "} + +data: {"type":"content","content":"the information I found, "} + +data: {"type":"content","content":"your application is healthy."} + +data: {"type":"usage","usage":{"input_tokens":150,"output_tokens":42,"total_tokens":192,"max_tokens":131072}} + +data: {"type":"sources","sources":[{"title":"Deploying Applications","url":"https://docs.nais.io/deploy/"}]} + +data: {"type":"done"} +``` + +#### Example: Response with thinking + +When thinking mode is enabled, reasoning chunks arrive before and between tool calls: + +``` +data: {"type":"metadata","conversation_id":"550e8400-e29b-41d4-a716-446655440000","message_id":"6ba7b810-9dad-11d1-80b4-00c04fd430c8"} + +data: {"type":"thinking","thinking":"The user wants to know about deployment. I should check their applications first."} + +data: {"type":"tool_start","tool_call_id":"call_123","tool_name":"execute_graphql","description":"Executing execute_graphql..."} + +data: {"type":"tool_end","tool_call_id":"call_123","tool_name":"execute_graphql","description":"Executed execute_graphql","tool_success":true} + +data: {"type":"thinking","thinking":"I have the application list. Now I can give deployment instructions."} + +data: {"type":"content","content":"To deploy your application, "} + +data: {"type":"content","content":"push to the main branch and the pipeline will handle the rest."} + +data: {"type":"usage","usage":{"input_tokens":200,"output_tokens":55,"total_tokens":255,"max_tokens":131072}} + +data: {"type":"sources","sources":[{"title":"CI/CD Guide","url":"https://docs.nais.io/cicd/"}]} + +data: {"type":"done"} +``` + +#### Example: Response with a chart + +``` +data: {"type":"metadata","conversation_id":"550e8400-e29b-41d4-a716-446655440000","message_id":"6ba7b810-9dad-11d1-80b4-00c04fd430c8"} + +data: {"type":"tool_start","tool_call_id":"call_456","tool_name":"render_chart","description":"Executing render_chart..."} + +data: {"type":"tool_end","tool_call_id":"call_456","tool_name":"render_chart","description":"Executed render_chart","tool_success":true} + +data: {"type":"chart","chart":{"chart_type":"line","title":"CPU Usage for my-app","environment":"dev","query":"sum(rate(container_cpu_usage_seconds_total{app=\"my-app\"}[5m])) by (pod)","interval":"1h","y_format":"cpu_cores","label_template":"{pod}"}} + +data: {"type":"content","content":"Here's your CPU usage over the last hour. "} + +data: {"type":"content","content":"Usage looks stable around 0.5 cores."} + +data: {"type":"usage","usage":{"input_tokens":180,"output_tokens":38,"total_tokens":218,"max_tokens":131072}} + +data: {"type":"done"} +``` + +--- + +### GET /agent/conversations + +List all conversations for the authenticated user, ordered by most recently updated. + +**Response:** +```json +{ + "conversations": [ + { + "id": "550e8400-e29b-41d4-a716-446655440000", + "title": "How do I deploy my application?", + "updated_at": "2024-01-15T10:30:00Z" + } + ] +} +``` + +--- + +### GET /agent/conversations/{conversationID} + +Get a specific conversation with all its messages. + +**Response:** +```json +{ + "conversation": { + "id": "550e8400-e29b-41d4-a716-446655440000", + "title": "How do I deploy my application?", + "messages": [ + { + "id": "6ba7b810-9dad-11d1-80b4-00c04fd430c8", + "role": "user", + "content": "How do I deploy my application?", + "created_at": "2024-01-15T10:30:00Z" + }, + { + "id": "6ba7b811-9dad-11d1-80b4-00c04fd430c8", + "role": "assistant", + "content": "To deploy your application...", + "blocks": [ + { + "type": "thinking", + "thinking": "The user wants to deploy. Let me check their pipeline setup." + }, + { + "type": "text", + "text": "To deploy your application..." + } + ], + "sources": [ + { + "title": "Deploying Applications", + "url": "https://docs.nais.io/deploy/" + } + ], + "created_at": "2024-01-15T10:30:05Z" + } + ], + "created_at": "2024-01-15T10:30:00Z", + "updated_at": "2024-01-15T10:30:05Z" + } +} +``` + +#### Content Blocks + +Assistant messages contain a `blocks` array representing the content in the order it was produced. Render them top-to-bottom to reconstruct the full response. + +**Block types:** + +| Type | Fields | Description | +|------|--------|-------------| +| `thinking` | `thinking` | Model's reasoning/thought process | +| `text` | `text` | Regular text output | +| `chart` | `chart` | A Prometheus chart to render — same schema as the SSE `chart` event | + +Tool calls and token usage are internal details and are not included in the blocks returned to clients. + +**Example with all block types:** +```json +{ + "blocks": [ + { + "type": "thinking", + "thinking": "The user is asking about CPU usage. I should visualize this." + }, + { + "type": "text", + "text": "Here's the CPU usage for your application:" + }, + { + "type": "chart", + "chart": { + "chart_type": "line", + "title": "CPU Usage for my-app", + "environment": "dev", + "query": "sum(rate(container_cpu_usage_seconds_total{app=\"my-app\"}[5m])) by (pod)", + "interval": "1h", + "y_format": "cpu_cores", + "label_template": "{pod}" + } + }, + { + "type": "text", + "text": "Usage has been stable around 0.5 cores." + } + ] +} +``` + +--- + +### DELETE /agent/conversations/{conversationID} + +Delete a conversation and all its messages. + +**Response:** +```json +{ + "deleted": true +} +``` + +--- + +## Error Responses + +All endpoints return errors in the following format: + +```json +{ + "error": "Description of the error" +} +``` + +**HTTP Status Codes:** +- `400 Bad Request` — Invalid request body or parameters +- `401 Unauthorized` — Authentication required +- `404 Not Found` — Conversation not found +- `500 Internal Server Error` — Server error + +--- + +## Client Implementation Notes + +### Handling SSE Streams + +1. Set appropriate request headers: + ``` + Accept: text/event-stream + Cache-Control: no-cache + ``` +2. Parse each `data:` line as JSON. +3. Accumulate `content` events to build the full response text. +4. Use `tool_start` / `tool_end` events to show a progress indicator. +5. Render `chart` events inline as Prometheus chart components. +6. Optionally display `thinking` events as a collapsible reasoning section. +7. Stop processing on `done` or `error`. + +### Building the Display from Stream Events + +To reconstruct the same block layout as returned by `GET /conversations/{id}`, track state as events arrive: + +1. On `thinking`: append a `thinking` block with the accumulated text. +2. On `tool_start` / `tool_end`: show a transient tool indicator (not stored as a block). +3. On `chart`: append a `chart` block. +4. On `content`: buffer the text; flush as a `text` block when the stream ends or a `chart` / `tool_start` interrupts it. +5. On `usage` / `sources` / `done`: finalise the display. + +### Thinking Mode + +When thinking mode is enabled on the server (`AGENT_VERTEX_AI_INCLUDE_THOUGHTS=true`), `thinking` events are streamed before the model's response. Thinking blocks are also persisted and returned by `GET /conversations/{id}`. + +Consider: +- Showing a "thinking…" animation while `thinking` events are arriving. +- Rendering thinking blocks as a collapsible section in the conversation history. +- Providing a UI toggle so users can show or hide thinking blocks. + +Thinking events are only emitted when using a Gemini model with thinking mode enabled. If not enabled, no `thinking` events are sent. + +### Conversation Management + +- A new conversation is created automatically when `conversation_id` is omitted. +- Conversation titles are derived from the first message. +- There is no hard cap on the number of conversations per user. + +### Sources + +Every response includes a `sources` array containing the documentation pages that were retrieved as context for the query (up to 5). These are always the pages the model had available — not a filtered subset — so the client can display all of them as "references used". + +In the streaming endpoint, sources arrive as a single `sources` event just before `done`. In the non-streaming endpoint and in `GET /conversations/{id}`, they are included directly in the response / message object. + +### Rendering Charts + +When the model determines that metrics data is better visualised as a chart it emits a `chart` event (streaming) or a `chart` block (in stored conversations). Render a Prometheus chart component using the provided data. + +#### Chart Schema + +| Field | Type | Required | Description | +|-------|------|----------|-------------| +| `chart_type` | `"line"` | Yes | Chart type — only `line` is currently supported | +| `title` | string | Yes | Human-readable title | +| `environment` | string | Yes | Environment to query metrics from (e.g. `dev`, `prod`) | +| `query` | string | Yes | PromQL query to execute | +| `interval` | string | No | Time window: `1h`, `6h`, `1d`, `7d`, or `30d` (default: `1h`) | +| `y_format` | string | No | Y-axis format: `number`, `percentage`, `bytes`, `cpu_cores`, or `duration` | +| `label_template` | string | No | Series label template using `{label_name}` syntax, e.g. `{pod}` or `{pod}/{container}` | + +#### Mapping to a chart component + +```typescript +type PrometheusChartProps = { + environmentName: string; // chart.environment + query: string; // chart.query + interval?: string; // chart.interval (default "1h") + labelFormatter: (labels: { name: string; value: string }[]) => string; + formatYValue?: (value: number) => string; +}; +``` + +#### Label Template Processing + +```typescript +function createLabelFormatter(template: string | undefined) { + if (!template) { + return (labels: { name: string; value: string }[]) => + labels.map(l => l.value).join('/'); + } + return (labels: { name: string; value: string }[]) => { + let result = template; + for (const label of labels) { + result = result.replace(`{${label.name}}`, label.value); + } + return result; + }; +} +``` + +#### Y-Axis Format + +| `y_format` | Description | Example | +|------------|-------------|---------| +| `number` | Plain number | `1,234.56` | +| `percentage` | Percentage | `45.2%` | +| `bytes` | Byte size | `1.5 GiB` | +| `cpu_cores` | CPU cores | `0.5 cores` | +| `duration` | Time duration | `2m 30s` | + +#### Interval Values + +| `interval` | Description | +|------------|-------------| +| `1h` | Last 1 hour | +| `6h` | Last 6 hours | +| `1d` | Last 1 day | +| `7d` | Last 7 days | +| `30d` | Last 30 days | \ No newline at end of file diff --git a/internal/agent/agentsql/agent.sql.go b/internal/agent/agentsql/agent.sql.go new file mode 100644 index 000000000..71dc23c7c --- /dev/null +++ b/internal/agent/agentsql/agent.sql.go @@ -0,0 +1,283 @@ +// Code generated by sqlc. DO NOT EDIT. +// source: agent.sql + +package agentsql + +import ( + "context" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgtype" +) + +const conversationExists = `-- name: ConversationExists :one +SELECT + EXISTS ( + SELECT + 1 + FROM + agent_conversations + WHERE + id = $1 + AND user_id = $2 + ) +` + +type ConversationExistsParams struct { + ID uuid.UUID + UserID uuid.UUID +} + +func (q *Queries) ConversationExists(ctx context.Context, arg ConversationExistsParams) (bool, error) { + row := q.db.QueryRow(ctx, conversationExists, arg.ID, arg.UserID) + var exists bool + err := row.Scan(&exists) + return exists, err +} + +const createConversation = `-- name: CreateConversation :one +INSERT INTO + agent_conversations (user_id, title) +VALUES + ($1, $2) +RETURNING + id +` + +type CreateConversationParams struct { + UserID uuid.UUID + Title string +} + +func (q *Queries) CreateConversation(ctx context.Context, arg CreateConversationParams) (uuid.UUID, error) { + row := q.db.QueryRow(ctx, createConversation, arg.UserID, arg.Title) + var id uuid.UUID + err := row.Scan(&id) + return id, err +} + +const deleteConversation = `-- name: DeleteConversation :execrows +DELETE FROM agent_conversations +WHERE + id = $1 + AND user_id = $2 +` + +type DeleteConversationParams struct { + ID uuid.UUID + UserID uuid.UUID +} + +func (q *Queries) DeleteConversation(ctx context.Context, arg DeleteConversationParams) (int64, error) { + result, err := q.db.Exec(ctx, deleteConversation, arg.ID, arg.UserID) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + +const getConversation = `-- name: GetConversation :one +SELECT + id, + title, + created_at, + updated_at +FROM + agent_conversations +WHERE + id = $1 + AND user_id = $2 +` + +type GetConversationParams struct { + ID uuid.UUID + UserID uuid.UUID +} + +type GetConversationRow struct { + ID uuid.UUID + Title string + CreatedAt pgtype.Timestamptz + UpdatedAt pgtype.Timestamptz +} + +func (q *Queries) GetConversation(ctx context.Context, arg GetConversationParams) (*GetConversationRow, error) { + row := q.db.QueryRow(ctx, getConversation, arg.ID, arg.UserID) + var i GetConversationRow + err := row.Scan( + &i.ID, + &i.Title, + &i.CreatedAt, + &i.UpdatedAt, + ) + return &i, err +} + +const getConversationHistory = `-- name: GetConversationHistory :many +SELECT + role, + content, + blocks +FROM + agent_messages +WHERE + conversation_id = $1 +ORDER BY + created_at ASC +` + +type GetConversationHistoryRow struct { + Role string + Content string + Blocks []byte +} + +func (q *Queries) GetConversationHistory(ctx context.Context, conversationID uuid.UUID) ([]*GetConversationHistoryRow, error) { + rows, err := q.db.Query(ctx, getConversationHistory, conversationID) + if err != nil { + return nil, err + } + defer rows.Close() + items := []*GetConversationHistoryRow{} + for rows.Next() { + var i GetConversationHistoryRow + if err := rows.Scan(&i.Role, &i.Content, &i.Blocks); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const getConversationMessages = `-- name: GetConversationMessages :many +SELECT + id, + role, + content, + sources, + blocks, + created_at +FROM + agent_messages +WHERE + conversation_id = $1 +ORDER BY + created_at ASC +` + +type GetConversationMessagesRow struct { + ID uuid.UUID + Role string + Content string + Sources []byte + Blocks []byte + CreatedAt pgtype.Timestamptz +} + +func (q *Queries) GetConversationMessages(ctx context.Context, conversationID uuid.UUID) ([]*GetConversationMessagesRow, error) { + rows, err := q.db.Query(ctx, getConversationMessages, conversationID) + if err != nil { + return nil, err + } + defer rows.Close() + items := []*GetConversationMessagesRow{} + for rows.Next() { + var i GetConversationMessagesRow + if err := rows.Scan( + &i.ID, + &i.Role, + &i.Content, + &i.Sources, + &i.Blocks, + &i.CreatedAt, + ); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const insertMessage = `-- name: InsertMessage :exec +INSERT INTO + agent_messages (conversation_id, role, content, sources, blocks) +VALUES + ($1, $2, $3, $4, $5) +` + +type InsertMessageParams struct { + ConversationID uuid.UUID + Role string + Content string + Sources []byte + Blocks []byte +} + +func (q *Queries) InsertMessage(ctx context.Context, arg InsertMessageParams) error { + _, err := q.db.Exec(ctx, insertMessage, + arg.ConversationID, + arg.Role, + arg.Content, + arg.Sources, + arg.Blocks, + ) + return err +} + +const listConversations = `-- name: ListConversations :many +SELECT + id, + title, + updated_at +FROM + agent_conversations +WHERE + user_id = $1 +ORDER BY + updated_at DESC +` + +type ListConversationsRow struct { + ID uuid.UUID + Title string + UpdatedAt pgtype.Timestamptz +} + +func (q *Queries) ListConversations(ctx context.Context, userID uuid.UUID) ([]*ListConversationsRow, error) { + rows, err := q.db.Query(ctx, listConversations, userID) + if err != nil { + return nil, err + } + defer rows.Close() + items := []*ListConversationsRow{} + for rows.Next() { + var i ListConversationsRow + if err := rows.Scan(&i.ID, &i.Title, &i.UpdatedAt); err != nil { + return nil, err + } + items = append(items, &i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const touchConversation = `-- name: TouchConversation :exec +UPDATE agent_conversations +SET + updated_at = NOW() +WHERE + id = $1 +` + +func (q *Queries) TouchConversation(ctx context.Context, id uuid.UUID) error { + _, err := q.db.Exec(ctx, touchConversation, id) + return err +} diff --git a/internal/agent/agentsql/db.go b/internal/agent/agentsql/db.go new file mode 100644 index 000000000..afa337f08 --- /dev/null +++ b/internal/agent/agentsql/db.go @@ -0,0 +1,30 @@ +// Code generated by sqlc. DO NOT EDIT. + +package agentsql + +import ( + "context" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" +) + +type DBTX interface { + Exec(context.Context, string, ...interface{}) (pgconn.CommandTag, error) + Query(context.Context, string, ...interface{}) (pgx.Rows, error) + QueryRow(context.Context, string, ...interface{}) pgx.Row +} + +func New(db DBTX) *Queries { + return &Queries{db: db} +} + +type Queries struct { + db DBTX +} + +func (q *Queries) WithTx(tx pgx.Tx) *Queries { + return &Queries{ + db: tx, + } +} diff --git a/internal/agent/agentsql/models.go b/internal/agent/agentsql/models.go new file mode 100644 index 000000000..24721ce58 --- /dev/null +++ b/internal/agent/agentsql/models.go @@ -0,0 +1,3 @@ +// Code generated by sqlc. DO NOT EDIT. + +package agentsql diff --git a/internal/agent/agentsql/querier.go b/internal/agent/agentsql/querier.go new file mode 100644 index 000000000..3e6aa6fe1 --- /dev/null +++ b/internal/agent/agentsql/querier.go @@ -0,0 +1,23 @@ +// Code generated by sqlc. DO NOT EDIT. + +package agentsql + +import ( + "context" + + "github.com/google/uuid" +) + +type Querier interface { + ConversationExists(ctx context.Context, arg ConversationExistsParams) (bool, error) + CreateConversation(ctx context.Context, arg CreateConversationParams) (uuid.UUID, error) + DeleteConversation(ctx context.Context, arg DeleteConversationParams) (int64, error) + GetConversation(ctx context.Context, arg GetConversationParams) (*GetConversationRow, error) + GetConversationHistory(ctx context.Context, conversationID uuid.UUID) ([]*GetConversationHistoryRow, error) + GetConversationMessages(ctx context.Context, conversationID uuid.UUID) ([]*GetConversationMessagesRow, error) + InsertMessage(ctx context.Context, arg InsertMessageParams) error + ListConversations(ctx context.Context, userID uuid.UUID) ([]*ListConversationsRow, error) + TouchConversation(ctx context.Context, id uuid.UUID) error +} + +var _ Querier = (*Queries)(nil) diff --git a/internal/agent/chat/chat.go b/internal/agent/chat/chat.go new file mode 100644 index 000000000..29afae99b --- /dev/null +++ b/internal/agent/chat/chat.go @@ -0,0 +1,630 @@ +// Package chat provides the LLM chat client for the agent using Vertex AI Gemini. +package chat + +import ( + "context" + "fmt" + + "github.com/nais/api/internal/agent/rag" + "github.com/sirupsen/logrus" + "google.golang.org/genai" +) + +// EmbeddingConfig holds configuration for the Vertex AI embedding client. +type EmbeddingConfig struct { + // ProjectID is the GCP project ID. + ProjectID string + + // Location is the GCP region (must be in EU, e.g., "europe-west1"). + Location string + + // ModelName is the embedding model to use (e.g., "gemini-embedding-001"). + ModelName string +} + +// EmbeddingClient generates embeddings using Vertex AI. +type EmbeddingClient struct { + client *genai.Client + model string + log logrus.FieldLogger +} + +// NewEmbeddingClient creates a new Vertex AI embedding client. +func NewEmbeddingClient(ctx context.Context, cfg EmbeddingConfig, log logrus.FieldLogger) (*EmbeddingClient, error) { + client, err := genai.NewClient(ctx, &genai.ClientConfig{ + Project: cfg.ProjectID, + Location: cfg.Location, + Backend: genai.BackendVertexAI, + }) + if err != nil { + return nil, fmt.Errorf("failed to create Vertex AI client: %w", err) + } + + log.WithFields(logrus.Fields{ + "project": cfg.ProjectID, + "location": cfg.Location, + "model": cfg.ModelName, + }).Info("initialized Vertex AI embedding client") + + return &EmbeddingClient{ + client: client, + model: cfg.ModelName, + log: log, + }, nil +} + +// Embed returns the embedding vector for the given text. +func (c *EmbeddingClient) Embed(ctx context.Context, text string) ([]float32, error) { + embeddings, err := c.EmbedBatch(ctx, []string{text}) + if err != nil { + return nil, err + } + return embeddings[0], nil +} + +// EmbedBatch returns embedding vectors for multiple texts in a single API call. +// This is much more efficient than calling Embed multiple times. +// The Vertex AI API supports up to 250 texts per batch request. +func (c *EmbeddingClient) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error) { + if len(texts) == 0 { + return nil, nil + } + + contents := make([]*genai.Content, len(texts)) + for i, text := range texts { + contents[i] = &genai.Content{Parts: []*genai.Part{{Text: text}}} + } + + result, err := c.client.Models.EmbedContent(ctx, c.model, contents, nil) + if err != nil { + return nil, fmt.Errorf("failed to embed texts: %w", err) + } + + if result == nil || len(result.Embeddings) != len(texts) { + return nil, fmt.Errorf("expected %d embeddings, got %d", len(texts), len(result.Embeddings)) + } + + embeddings := make([][]float32, len(texts)) + for i, emb := range result.Embeddings { + embeddings[i] = emb.Values + } + + return embeddings, nil +} + +// Close cleans up resources. +func (c *EmbeddingClient) Close() error { + // The genai client doesn't have a Close method + return nil +} + +const ( + defaultTemperature = 0.3 + defaultMaxTokens = 4096 + defaultTopP = 0.95 + defaultTopK = 40 +) + +// Role represents the role of a message sender. +type Role string + +const ( + RoleUser Role = "user" + RoleAssistant Role = "assistant" + RoleTool Role = "tool" +) + +// Client defines the interface for LLM chat interactions. +// This interface exists primarily for testing purposes. +type Client interface { + // Chat sends a message to the LLM and returns a response. + // The response may contain tool calls that need to be executed. + Chat(ctx context.Context, req *Request) (*Response, error) + + // Close cleans up any resources held by the client. + Close() error +} + +// StreamingClient extends Client with streaming support. +type StreamingClient interface { + Client + + // ChatStream sends a message and returns a channel of response chunks. + // The channel is closed when the response is complete or an error occurs. + ChatStream(ctx context.Context, req *Request) (<-chan StreamChunk, error) +} + +// Request represents a chat request to the LLM. +type Request struct { + // SystemPrompt is the system-level instruction for the LLM. + SystemPrompt string + + // Messages is the conversation history. + Messages []Message + + // Tools is the list of tools available to the LLM. + Tools []ToolDefinition + + // Documents contains RAG results to include in context. + Documents []rag.Document +} + +// Response represents a chat response from the LLM. +type Response struct { + // Content is the text response from the LLM. + Content string + + // Thinking is the model's reasoning/thought process (only populated when thinking mode is enabled). + Thinking string + + // ToolCalls contains any tool calls the LLM wants to make. + ToolCalls []ToolCall + + // Usage contains token usage statistics. + Usage *UsageStats +} + +// StreamChunk represents a chunk of a streaming response. +type StreamChunk struct { + // Content is the partial text content. + Content string + + // Thinking is the model's reasoning/thought process (only populated when thinking mode is enabled). + Thinking string + + // ToolCalls contains tool calls (sent as complete objects). + ToolCalls []ToolCall + + // Usage contains token usage statistics. + Usage *UsageStats + + // Done is true if this is the final chunk. + Done bool + + // Error is non-nil if an error occurred. + Error error +} + +// Message represents a single message in the conversation. +type Message struct { + // Role is the sender role (user, assistant, or tool). + Role Role + + // Content is the message text. + Content string + + // ToolCallID is set for tool response messages to indicate which tool call this responds to. + ToolCallID string + + // ToolCalls is set for assistant messages that request tool calls. + ToolCalls []ToolCall +} + +// ToolCall represents a tool invocation requested by the LLM. +type ToolCall struct { + // ID is a unique identifier for this tool call. + ID string + + // Name is the name of the tool to invoke. + Name string + + // Arguments contains the arguments for the tool. + Arguments map[string]any + + // ThoughtSignature is an encrypted representation of the model's internal thought process. + // Required by Gemini 3 models for function calling - must be preserved and sent back + // when providing function responses. + ThoughtSignature string `json:"thought_signature,omitempty"` +} + +// ToolDefinition describes a tool available to the LLM. +type ToolDefinition struct { + // Name is the unique name of the tool. + Name string + + // Description explains what the tool does. + Description string + + // Parameters describes the tool's parameters. + Parameters []ParameterDefinition +} + +// ParameterDefinition describes a single parameter for a tool. +type ParameterDefinition struct { + // Name is the parameter name. + Name string + + // Type is the parameter type (e.g., "string", "object", "array"). + Type string + + // Description explains what the parameter is for. + Description string + + // Required indicates if the parameter must be provided. + Required bool +} + +// UsageStats contains token usage information. +type UsageStats struct { + // InputTokens is the number of tokens in the input. + InputTokens int `json:"input_tokens"` + + // OutputTokens is the number of tokens in the output. + OutputTokens int `json:"output_tokens"` + + // TotalTokens is the total number of tokens used (input + output). + TotalTokens int `json:"total_tokens"` + + // MaxTokens is the maximum number of tokens allowed for the model/context window. + MaxTokens int `json:"max_tokens,omitempty"` +} + +// Config holds configuration for the Vertex AI chat client. +type Config struct { + ProjectID string + Location string + ModelName string + IncludeThoughts bool // Include thought content in responses +} + +// VertexAIClient implements StreamingClient using Vertex AI Gemini. +type VertexAIClient struct { + client *genai.Client + modelName string + log logrus.FieldLogger + includeThoughts bool +} + +// NewClient creates a new Vertex AI chat client. +func NewClient(ctx context.Context, cfg Config, log logrus.FieldLogger) (*VertexAIClient, error) { + client, err := genai.NewClient(ctx, &genai.ClientConfig{ + Project: cfg.ProjectID, + Location: cfg.Location, + Backend: genai.BackendVertexAI, + }) + if err != nil { + return nil, fmt.Errorf("failed to create Vertex AI client: %w", err) + } + + log.WithFields(logrus.Fields{ + "project": cfg.ProjectID, + "location": cfg.Location, + "model": cfg.ModelName, + }).Info("initialized Vertex AI chat client") + + return &VertexAIClient{ + client: client, + modelName: cfg.ModelName, + log: log, + includeThoughts: cfg.IncludeThoughts, + }, nil +} + +// Chat sends a message to the LLM and returns a response. +func (c *VertexAIClient) Chat(ctx context.Context, req *Request) (*Response, error) { + config := c.buildGenerateContentConfig(req) + contents := c.buildContents(req) + + resp, err := c.client.Models.GenerateContent(ctx, c.modelName, contents, config) + if err != nil { + return nil, fmt.Errorf("failed to generate content: %w", err) + } + + return c.convertResponse(resp), nil +} + +// ChatStream sends a message and returns a channel of response chunks. +func (c *VertexAIClient) ChatStream(ctx context.Context, req *Request) (<-chan StreamChunk, error) { + chunkCh := make(chan StreamChunk, 100) + + go func() { + defer close(chunkCh) + + config := c.buildGenerateContentConfig(req) + contents := c.buildContents(req) + + for resp, err := range c.client.Models.GenerateContentStream(ctx, c.modelName, contents, config) { + if err != nil { + chunkCh <- StreamChunk{Error: err} + return + } + + chunk := c.convertStreamResponse(resp) + chunkCh <- chunk + } + + chunkCh <- StreamChunk{Done: true} + }() + + return chunkCh, nil +} + +// Close cleans up resources. +func (c *VertexAIClient) Close() error { + // The genai client doesn't have a Close method + return nil +} + +func (c *VertexAIClient) buildGenerateContentConfig(req *Request) *genai.GenerateContentConfig { + temp := float32(defaultTemperature) + topP := float32(defaultTopP) + topK := float32(defaultTopK) + + config := &genai.GenerateContentConfig{ + Temperature: &temp, + MaxOutputTokens: int32(defaultMaxTokens), + TopP: &topP, + TopK: &topK, + Tools: c.convertTools(req.Tools), + } + + if req.SystemPrompt != "" { + config.SystemInstruction = &genai.Content{ + Parts: []*genai.Part{{Text: req.SystemPrompt}}, + } + } + + // Enable thinking mode for Gemini 3+ models + // This is required for proper thought signature handling + config.ThinkingConfig = &genai.ThinkingConfig{ + IncludeThoughts: c.includeThoughts, + } + + return config +} + +func (c *VertexAIClient) buildContents(req *Request) []*genai.Content { + contents := make([]*genai.Content, 0, len(req.Messages)) + + // Collect consecutive tool responses to batch them into a single content block. + // Vertex AI/Gemini requires all function responses for a turn to be in one content. + var pendingToolResponses []*genai.Part + + flushToolResponses := func() { + if len(pendingToolResponses) > 0 { + contents = append(contents, &genai.Content{ + Role: "user", + Parts: pendingToolResponses, + }) + pendingToolResponses = nil + } + } + + for _, msg := range req.Messages { + switch msg.Role { + case RoleTool: + // Collect tool response - will be batched with other consecutive tool responses + pendingToolResponses = append(pendingToolResponses, &genai.Part{ + FunctionResponse: &genai.FunctionResponse{ + Name: msg.ToolCallID, + Response: map[string]any{"result": msg.Content}, + }, + }) + case RoleAssistant: + // Flush any pending tool responses before adding assistant message + flushToolResponses() + + content := &genai.Content{ + Role: "model", + } + if len(msg.ToolCalls) > 0 { + // Assistant message with tool calls + parts := make([]*genai.Part, 0, len(msg.ToolCalls)+1) + if msg.Content != "" { + parts = append(parts, &genai.Part{Text: msg.Content}) + } + for _, tc := range msg.ToolCalls { + part := &genai.Part{ + FunctionCall: &genai.FunctionCall{ + Name: tc.Name, + Args: tc.Arguments, + }, + } + // Preserve thought signature for Gemini 3 models + if tc.ThoughtSignature != "" { + part.ThoughtSignature = []byte(tc.ThoughtSignature) + } + parts = append(parts, part) + } + content.Parts = parts + } else { + content.Parts = []*genai.Part{{Text: msg.Content}} + } + contents = append(contents, content) + default: + // Flush any pending tool responses before adding user message + flushToolResponses() + + content := &genai.Content{ + Role: "user", + Parts: []*genai.Part{{Text: msg.Content}}, + } + contents = append(contents, content) + } + } + + // Flush any remaining tool responses at the end + flushToolResponses() + + // Add document context if available + if len(req.Documents) > 0 { + docContext := formatDocuments(req.Documents) + if len(contents) > 0 && contents[len(contents)-1].Role == "user" { + // Prepend to last user message + lastContent := contents[len(contents)-1] + newParts := make([]*genai.Part, 0, len(lastContent.Parts)+1) + newParts = append(newParts, &genai.Part{Text: docContext + "\n\n"}) + newParts = append(newParts, lastContent.Parts...) + lastContent.Parts = newParts + } + } + + return contents +} + +func formatDocuments(docs []rag.Document) string { + if len(docs) == 0 { + return "" + } + + result := "Here is relevant documentation:\n\n" + for _, doc := range docs { + result += fmt.Sprintf("### %s\n%s\nSource: %s\n\n", doc.Title, doc.Content, doc.URL) + } + return result +} + +func (c *VertexAIClient) convertTools(tools []ToolDefinition) []*genai.Tool { + if len(tools) == 0 { + return nil + } + + funcDecls := make([]*genai.FunctionDeclaration, 0, len(tools)) + for _, tool := range tools { + funcDecl := &genai.FunctionDeclaration{ + Name: tool.Name, + Description: tool.Description, + Parameters: c.convertParameters(tool.Parameters), + } + funcDecls = append(funcDecls, funcDecl) + } + + return []*genai.Tool{ + {FunctionDeclarations: funcDecls}, + } +} + +func (c *VertexAIClient) convertParameters(params []ParameterDefinition) *genai.Schema { + if len(params) == 0 { + return nil + } + + properties := make(map[string]*genai.Schema) + required := make([]string, 0) + + for _, param := range params { + properties[param.Name] = &genai.Schema{ + Type: convertType(param.Type), + Description: param.Description, + } + if param.Required { + required = append(required, param.Name) + } + } + + return &genai.Schema{ + Type: genai.TypeObject, + Properties: properties, + Required: required, + } +} + +func convertType(t string) genai.Type { + switch t { + case "string": + return genai.TypeString + case "number": + return genai.TypeNumber + case "integer": + return genai.TypeInteger + case "boolean": + return genai.TypeBoolean + case "array": + return genai.TypeArray + case "object": + return genai.TypeObject + default: + return genai.TypeString + } +} + +func (c *VertexAIClient) convertResponse(resp *genai.GenerateContentResponse) *Response { + result := &Response{ + Usage: &UsageStats{}, + } + + if resp.UsageMetadata != nil { + result.Usage.InputTokens = int(resp.UsageMetadata.PromptTokenCount) + result.Usage.OutputTokens = int(resp.UsageMetadata.CandidatesTokenCount) + result.Usage.TotalTokens = int(resp.UsageMetadata.PromptTokenCount + resp.UsageMetadata.CandidatesTokenCount) + result.Usage.MaxTokens = c.getMaxContextWindow() + } + + if len(resp.Candidates) == 0 { + return result + } + + candidate := resp.Candidates[0] + if candidate.Content == nil { + return result + } + + for _, part := range candidate.Content.Parts { + if part.Text != "" { + // Check if this is a thought part (model's reasoning) + if part.Thought { + result.Thinking += part.Text + } else { + result.Content += part.Text + } + } + if part.FunctionCall != nil { + result.ToolCalls = append(result.ToolCalls, ToolCall{ + ID: part.FunctionCall.Name, // Use name as ID for Gemini + Name: part.FunctionCall.Name, + Arguments: part.FunctionCall.Args, + ThoughtSignature: string(part.ThoughtSignature), // Preserve thought signature for Gemini 3 + }) + } + } + + return result +} + +func (c *VertexAIClient) convertStreamResponse(resp *genai.GenerateContentResponse) StreamChunk { + chunk := StreamChunk{} + + if resp.UsageMetadata != nil { + chunk.Usage = &UsageStats{ + InputTokens: int(resp.UsageMetadata.PromptTokenCount), + OutputTokens: int(resp.UsageMetadata.CandidatesTokenCount), + TotalTokens: int(resp.UsageMetadata.PromptTokenCount + resp.UsageMetadata.CandidatesTokenCount), + MaxTokens: c.getMaxContextWindow(), + } + } + + if len(resp.Candidates) == 0 { + return chunk + } + + candidate := resp.Candidates[0] + if candidate.Content == nil { + return chunk + } + + for _, part := range candidate.Content.Parts { + if part.Text != "" { + // Check if this is a thought part (model's reasoning) + if part.Thought { + chunk.Thinking += part.Text + } else { + chunk.Content += part.Text + } + } + if part.FunctionCall != nil { + chunk.ToolCalls = append(chunk.ToolCalls, ToolCall{ + ID: part.FunctionCall.Name, + Name: part.FunctionCall.Name, + Arguments: part.FunctionCall.Args, + ThoughtSignature: string(part.ThoughtSignature), // Preserve thought signature for Gemini 3 + }) + } + } + + return chunk +} + +func (c *VertexAIClient) getMaxContextWindow() int { + // Default to 128k as a safe middle ground for modern models if unknown + return 131072 +} diff --git a/internal/agent/dataloader.go b/internal/agent/dataloader.go new file mode 100644 index 000000000..cf45f5c1b --- /dev/null +++ b/internal/agent/dataloader.go @@ -0,0 +1,39 @@ +package agent + +import ( + "context" + + "github.com/jackc/pgx/v5/pgxpool" + "github.com/nais/api/internal/agent/agentsql" + "github.com/nais/api/internal/database" +) + +type ctxKey int + +const loadersKey ctxKey = iota + +type loaders struct { + querier *agentsql.Queries +} + +// NewLoaderContext stores agent-specific loaders in the context so that package-level +// query functions can access the database without passing a store explicitly. +func NewLoaderContext(ctx context.Context, pool *pgxpool.Pool) context.Context { + return context.WithValue(ctx, loadersKey, &loaders{ + querier: agentsql.New(pool), + }) +} + +func fromContext(ctx context.Context) *loaders { + return ctx.Value(loadersKey).(*loaders) +} + +// db returns a transaction-aware querier. If a transaction is active in ctx +// (set by database.Transaction), it will be used; otherwise the base querier is returned. +func db(ctx context.Context) *agentsql.Queries { + l := fromContext(ctx) + if tx := database.TransactionFromContext(ctx); tx != nil { + return l.querier.WithTx(tx) + } + return l.querier +} diff --git a/internal/agent/handler.go b/internal/agent/handler.go new file mode 100644 index 000000000..38f3a6475 --- /dev/null +++ b/internal/agent/handler.go @@ -0,0 +1,466 @@ +// Package agent provides the AI chat service for the Nais platform. +package agent + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "strings" + + "github.com/99designs/gqlgen/graphql/handler" + "github.com/go-chi/chi/v5" + "github.com/google/uuid" + "github.com/nais/api/internal/agent/chat" + "github.com/nais/api/internal/agent/rag" + "github.com/nais/api/internal/agent/tools" + "github.com/nais/api/internal/auth/authz" + "github.com/nais/api/internal/graph/gengql" + "github.com/sirupsen/logrus" +) + +const maxRAGResults = 5 + +// Handler implements the HTTP handler for the agent chat service. +type Handler struct { + chatClient chat.StreamingClient + ragClient rag.DocumentSearcher + registry *tools.Registry + log logrus.FieldLogger +} + +// Config holds configuration for the agent handler. +type Config struct { + ChatClient chat.StreamingClient + RAGClient rag.DocumentSearcher + GraphQLHandler *handler.Server + TenantName string + Log logrus.FieldLogger +} + +// NewHandler creates a new agent HTTP handler. +func NewHandler(cfg Config) (*Handler, error) { + if cfg.GraphQLHandler == nil { + return nil, fmt.Errorf("GraphQLHandler is required") + } + + internalClient := NewInternalClient( + cfg.GraphQLHandler, + cfg.Log.WithField("component", "internal_client"), + ) + + schema := gengql.NewExecutableSchema(gengql.Config{}).Schema() + consoleBaseURL, urlPatterns := buildConsoleURLs(cfg.TenantName) + + registry := tools.NewRegistry(tools.RegistryConfig{ + Client: internalClient, + Schema: schema, + ConsoleBaseURL: consoleBaseURL, + ConsoleURLPatterns: urlPatterns, + }) + + return &Handler{ + chatClient: cfg.ChatClient, + ragClient: cfg.RAGClient, + registry: registry, + log: cfg.Log, + }, nil +} + +// RegisterRoutes registers the agent routes on the given router. +func (h *Handler) RegisterRoutes(r chi.Router) { + r.Post("/chat", h.Chat) + r.Post("/chat/stream", h.ChatStream) + r.Get("/conversations", h.ListConversations) + r.Get("/conversations/{conversationID}", h.GetConversation) + r.Delete("/conversations/{conversationID}", h.DeleteConversation) +} + +// ChatRequest represents a chat request. +type ChatRequest struct { + Message string `json:"message"` + ConversationID string `json:"conversation_id,omitempty"` + Context *ChatContext `json:"context,omitempty"` +} + +// ChatResponse represents a non-streaming chat response. +type ChatResponse struct { + ConversationID string `json:"conversation_id"` + MessageID string `json:"message_id"` + Content string `json:"content"` + Blocks []ContentBlock `json:"blocks,omitempty"` + Sources []Source `json:"sources,omitempty"` + Usage *chat.UsageStats `json:"usage,omitempty"` +} + +// StreamEvent represents a server-sent event for streaming responses. +type StreamEvent struct { + Type string `json:"type"` + ConversationID string `json:"conversation_id,omitempty"` + MessageID string `json:"message_id,omitempty"` + Content string `json:"content,omitempty"` + Thinking string `json:"thinking,omitempty"` + ToolName string `json:"tool_name,omitempty"` + ToolSuccess bool `json:"tool_success,omitempty"` + Description string `json:"description,omitempty"` + Sources []Source `json:"sources,omitempty"` + Chart *ChartData `json:"chart,omitempty"` + Usage *chat.UsageStats `json:"usage,omitempty"` + ToolCallID string `json:"tool_call_id,omitempty"` + ErrorCode string `json:"error_code,omitempty"` + ErrorMessage string `json:"error_message,omitempty"` +} + +// chatSetup holds the prepared state for a conversation turn, assembled by +// prepareConversation before being handed off to Chat or ChatStream. +type chatSetup struct { + Request ChatRequest + ConversationID uuid.UUID + History []chat.Message + Docs []rag.Document + Sources []Source +} + +// prepareConversation handles the common setup for chat requests: +// parses the request, authenticates the user, loads conversation history, +// and performs a RAG search for relevant documentation. +func (h *Handler) prepareConversation(w http.ResponseWriter, r *http.Request, log logrus.FieldLogger) (*chatSetup, error) { + ctx := r.Context() + + var req ChatRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + log.WithError(err).Debug("failed to decode request body") + writeJSONError(w, http.StatusBadRequest, "invalid request body") + return nil, err + } + + if strings.TrimSpace(req.Message) == "" { + writeJSONError(w, http.StatusBadRequest, "message is required") + return nil, fmt.Errorf("empty message") + } + + userID, err := getUserIDFromContext(ctx) + if err != nil { + log.WithError(err).Debug("failed to get user from context") + writeJSONError(w, http.StatusUnauthorized, "authentication required") + return nil, err + } + + conversationID, err := GetOrCreateConversation(ctx, userID, req.ConversationID, req.Message) + if err != nil { + log.WithError(err).Error("failed to get or create conversation") + writeJSONError(w, http.StatusInternalServerError, "failed to process conversation") + return nil, err + } + + history, err := GetConversationHistory(ctx, conversationID) + if err != nil { + log.WithError(err).Warn("failed to load conversation history, continuing without") + history = nil + } + + docs, sources, err := h.searchDocumentation(ctx, req.Message) + if err != nil { + log.WithError(err).Warn("RAG search failed, continuing without docs") + } + + return &chatSetup{ + Request: req, + ConversationID: conversationID, + History: history, + Docs: docs, + Sources: sources, + }, nil +} + +// Chat handles non-streaming chat requests. +func (h *Handler) Chat(w http.ResponseWriter, r *http.Request) { + ctx := r.Context() + log := h.log.WithField("method", "Chat") + + setup, err := h.prepareConversation(w, r, log) + if err != nil { + return + } + + log = log.WithField("conversation_id", setup.ConversationID) + + orchestrator := NewOrchestrator(h.chatClient, h.registry, log) + + result, err := orchestrator.Run(ctx, setup.Request.Message, setup.Request.Context, setup.Docs, setup.History) + if err != nil { + log.WithError(err).Error("orchestrator failed") + writeJSONError(w, http.StatusInternalServerError, "failed to process chat request") + return + } + + if err := StoreMessages(ctx, setup.ConversationID, setup.Request.Message, result, setup.Sources); err != nil { + log.WithError(err).Error("failed to store messages") + } + + textContent := extractTextContentFromBlocks(result.Blocks) + + writeJSON(w, http.StatusOK, ChatResponse{ + ConversationID: setup.ConversationID.String(), + MessageID: uuid.New().String(), + Content: textContent, + Blocks: clientVisibleBlocks(result.Blocks), + Sources: setup.Sources, + Usage: result.Usage, + }) +} + +// ChatStream handles streaming chat requests using Server-Sent Events. +func (h *Handler) ChatStream(w http.ResponseWriter, r *http.Request) { + ctx := r.Context() + log := h.log.WithField("method", "ChatStream") + + setup, err := h.prepareConversation(w, r, log) + if err != nil { + return + } + + log = log.WithField("conversation_id", setup.ConversationID) + + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.Header().Set("Connection", "keep-alive") + w.Header().Set("X-Accel-Buffering", "no") + + flusher, ok := w.(http.Flusher) + if !ok { + log.Error("response writer does not support flushing") + writeJSONError(w, http.StatusInternalServerError, "streaming not supported") + return + } + + messageID := uuid.New().String() + sendSSE(w, flusher, StreamEvent{ + Type: "metadata", + ConversationID: setup.ConversationID.String(), + MessageID: messageID, + }) + + orchestrator := NewOrchestrator(h.chatClient, h.registry, log) + streamCh := orchestrator.RunStream(ctx, setup.Request.Message, setup.Request.Context, setup.Docs, setup.History) + + for event := range streamCh { + switch event.Type { + case OrchestratorEventError: + log.WithError(event.Error).Error("stream error") + sendSSE(w, flusher, StreamEvent{ + Type: "error", + ErrorCode: "stream_error", + ErrorMessage: event.Error.Error(), + }) + return + + case OrchestratorEventDone: + result := event.Result + + if result.Usage != nil { + sendSSE(w, flusher, StreamEvent{ + Type: "usage", + Usage: result.Usage, + }) + } + + if len(setup.Sources) > 0 { + sendSSE(w, flusher, StreamEvent{ + Type: "sources", + Sources: setup.Sources, + }) + } + + sendSSE(w, flusher, StreamEvent{Type: "done"}) + + // Store messages asynchronously after the SSE stream is flushed. + // context.WithoutCancel preserves all context values (loaders, db pool, etc.) + // while detaching from the request's cancellation signal. + storeCtx := context.WithoutCancel(ctx) + go func(r *OrchestratorResult) { + if err := StoreMessages(storeCtx, setup.ConversationID, setup.Request.Message, r, setup.Sources); err != nil { + log.WithError(err).Error("failed to store messages") + } + }(result) + + return + + default: + sendSSE(w, flusher, event.ToStreamEvent()) + } + } + + // If the channel closed without a Done event (unexpected), send a done anyway. + sendSSE(w, flusher, StreamEvent{Type: "done"}) +} + +// ListConversations returns all conversations for the authenticated user. +func (h *Handler) ListConversations(w http.ResponseWriter, r *http.Request) { + ctx := r.Context() + + userID, err := getUserIDFromContext(ctx) + if err != nil { + writeJSONError(w, http.StatusUnauthorized, "authentication required") + return + } + + conversations, err := ListConversations(ctx, userID) + if err != nil { + h.log.WithError(err).Error("failed to list conversations") + writeJSONError(w, http.StatusInternalServerError, "failed to list conversations") + return + } + + writeJSON(w, http.StatusOK, map[string]any{ + "conversations": conversations, + }) +} + +// GetConversation returns a specific conversation with all its messages. +func (h *Handler) GetConversation(w http.ResponseWriter, r *http.Request) { + ctx := r.Context() + + userID, err := getUserIDFromContext(ctx) + if err != nil { + writeJSONError(w, http.StatusUnauthorized, "authentication required") + return + } + + conversationIDStr := chi.URLParam(r, "conversationID") + if conversationIDStr == "" { + writeJSONError(w, http.StatusBadRequest, "conversation_id is required") + return + } + + conversationID, err := uuid.Parse(conversationIDStr) + if err != nil { + writeJSONError(w, http.StatusBadRequest, "invalid conversation_id") + return + } + + conv, err := GetConversation(ctx, userID, conversationID) + if err != nil { + if errors.Is(err, ErrConversationNotFound) { + writeJSONError(w, http.StatusNotFound, "conversation not found") + return + } + h.log.WithError(err).Error("failed to get conversation") + writeJSONError(w, http.StatusInternalServerError, "failed to get conversation") + return + } + + writeJSON(w, http.StatusOK, map[string]any{ + "conversation": conv, + }) +} + +// DeleteConversation deletes a conversation. +func (h *Handler) DeleteConversation(w http.ResponseWriter, r *http.Request) { + ctx := r.Context() + + userID, err := getUserIDFromContext(ctx) + if err != nil { + writeJSONError(w, http.StatusUnauthorized, "authentication required") + return + } + + conversationIDStr := chi.URLParam(r, "conversationID") + if conversationIDStr == "" { + writeJSONError(w, http.StatusBadRequest, "conversation_id is required") + return + } + + conversationID, err := uuid.Parse(conversationIDStr) + if err != nil { + writeJSONError(w, http.StatusBadRequest, "invalid conversation_id") + return + } + + if err := DeleteConversation(ctx, userID, conversationID); err != nil { + if errors.Is(err, ErrConversationNotFound) { + writeJSONError(w, http.StatusNotFound, "conversation not found") + return + } + h.log.WithError(err).Error("failed to delete conversation") + writeJSONError(w, http.StatusInternalServerError, "failed to delete conversation") + return + } + + writeJSON(w, http.StatusOK, map[string]any{ + "deleted": true, + }) +} + +func (h *Handler) searchDocumentation(ctx context.Context, query string) ([]rag.Document, []Source, error) { + result, err := h.ragClient.Search(ctx, query, &rag.SearchOptions{ + MaxResults: maxRAGResults, + }) + if err != nil { + return nil, nil, err + } + + // Deduplicate sources by URL (multiple chunks from the same page → one source entry). + seen := make(map[string]bool) + sources := make([]Source, 0, len(result.Documents)) + for _, doc := range result.Documents { + if seen[doc.URL] { + continue + } + seen[doc.URL] = true + sources = append(sources, Source{ + Title: doc.Title, + URL: doc.URL, + }) + } + + return result.Documents, sources, nil +} + +// getUserIDFromContext extracts the user ID from the HTTP request context. +func getUserIDFromContext(ctx context.Context) (uuid.UUID, error) { + actor := authz.ActorFromContext(ctx) + if actor == nil || actor.User == nil { + return uuid.Nil, fmt.Errorf("no user in context") + } + return actor.User.GetID(), nil +} + +// buildConsoleURLs builds the console base URL and URL patterns for a tenant. +func buildConsoleURLs(tenantName string) (string, map[string]string) { + baseURL := fmt.Sprintf("https://console.%s.cloud.nais.io", tenantName) + + patterns := map[string]string{ + "team": baseURL + "/team/{team}", + "app": baseURL + "/team/{team}/{env}/app/{app}", + "job": baseURL + "/team/{team}/{env}/job/{job}", + "deployment": baseURL + "/team/{team}/deployments", + "cost": baseURL + "/team/{team}/cost", + "utilization": baseURL + "/team/{team}/utilization", + "secrets": baseURL + "/team/{team}/{env}/secret/{secret}", + "postgres": baseURL + "/team/{team}/{env}/postgres/{instance}", + "bucket": baseURL + "/team/{team}/{env}/bucket/{bucket}", + "redis": baseURL + "/team/{team}/{env}/redis/{instance}", + "opensearch": baseURL + "/team/{team}/{env}/opensearch/{instance}", + "kafka": baseURL + "/team/{team}/{env}/kafka/{topic}", + } + + return baseURL, patterns +} + +func sendSSE(w http.ResponseWriter, flusher http.Flusher, event StreamEvent) { + data, _ := json.Marshal(event) + fmt.Fprintf(w, "data: %s\n\n", data) + flusher.Flush() +} + +func writeJSON(w http.ResponseWriter, status int, v any) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + json.NewEncoder(w).Encode(v) +} + +func writeJSONError(w http.ResponseWriter, status int, message string) { + writeJSON(w, status, map[string]string{"error": message}) +} diff --git a/internal/agent/handler_test.go b/internal/agent/handler_test.go new file mode 100644 index 000000000..a3404b330 --- /dev/null +++ b/internal/agent/handler_test.go @@ -0,0 +1,55 @@ +package agent + +import ( + "encoding/json" + "testing" +) + +func TestContentBlockJSONRoundTrip(t *testing.T) { + // Test that ToolResult is properly serialized/deserialized in blocks + blocks := []ContentBlock{ + {Type: ContentBlockTypeText, Text: "Checking data..."}, + { + Type: ContentBlockTypeToolUse, + ToolCallID: "call_abc123", + ToolName: "execute_graphql", + ToolSuccess: true, + ToolResult: `{"data":{"team":{"slug":"my-team"}}}`, + }, + {Type: ContentBlockTypeText, Text: "Found your team!"}, + } + + // Marshal to JSON + jsonData, err := json.Marshal(blocks) + if err != nil { + t.Fatalf("failed to marshal blocks: %v", err) + } + + // Unmarshal back + var restored []ContentBlock + if err := json.Unmarshal(jsonData, &restored); err != nil { + t.Fatalf("failed to unmarshal blocks: %v", err) + } + + // Verify + if len(restored) != len(blocks) { + t.Fatalf("got %d blocks, want %d", len(restored), len(blocks)) + } + + toolBlock := restored[1] + if toolBlock.Type != ContentBlockTypeToolUse { + t.Errorf("block type = %v, want %v", toolBlock.Type, ContentBlockTypeToolUse) + } + if toolBlock.ToolCallID != "call_abc123" { + t.Errorf("ToolCallID = %q, want %q", toolBlock.ToolCallID, "call_abc123") + } + if toolBlock.ToolName != "execute_graphql" { + t.Errorf("ToolName = %q, want %q", toolBlock.ToolName, "execute_graphql") + } + if !toolBlock.ToolSuccess { + t.Error("ToolSuccess = false, want true") + } + if toolBlock.ToolResult != `{"data":{"team":{"slug":"my-team"}}}` { + t.Errorf("ToolResult = %q, want JSON data", toolBlock.ToolResult) + } +} diff --git a/internal/agent/internal_client.go b/internal/agent/internal_client.go new file mode 100644 index 000000000..7ac09f394 --- /dev/null +++ b/internal/agent/internal_client.go @@ -0,0 +1,168 @@ +package agent + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + + "github.com/99designs/gqlgen/graphql/handler" + "github.com/nais/api/internal/agent/tools" + "github.com/nais/api/internal/auth/authz" + "github.com/nais/api/internal/graph/pagination" + "github.com/nais/api/internal/team" + "github.com/sirupsen/logrus" +) + +// InternalClient executes GraphQL queries directly against the internal handler +// without going through a real HTTP connection. This is used by the tool integration +// to query the Nais API on behalf of users. +// +// InternalClient implements tools.GraphQLClient directly. +type InternalClient struct { + handler *handler.Server + log logrus.FieldLogger +} + +// Ensure InternalClient implements tools.GraphQLClient. +var _ tools.GraphQLClient = (*InternalClient)(nil) + +// NewInternalClient creates a new InternalClient for executing GraphQL queries +// directly against the internal handler. +func NewInternalClient(h *handler.Server, log logrus.FieldLogger) *InternalClient { + if log == nil { + log = logrus.StandardLogger() + } + return &InternalClient{ + handler: h, + log: log, + } +} + +// ExecuteGraphQL runs a GraphQL query with the given variables. +// Returns the data portion of the response as a map. +func (c *InternalClient) ExecuteGraphQL(ctx context.Context, query string, variables map[string]any) (map[string]any, error) { + c.log.WithFields(logrus.Fields{ + "query_length": len(query), + "has_variables": variables != nil, + }).Debug("executing internal GraphQL query") + + requestBody := map[string]any{ + "query": query, + } + if variables != nil { + requestBody["variables"] = variables + } + + body, err := json.Marshal(requestBody) + if err != nil { + return nil, fmt.Errorf("failed to marshal GraphQL request: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, "/graphql", bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + + rec := httptest.NewRecorder() + c.handler.ServeHTTP(rec, req) + + resp := rec.Result() + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("failed to read response: %w", err) + } + + c.log.WithFields(logrus.Fields{ + "status_code": resp.StatusCode, + "response_length": len(respBody), + }).Debug("GraphQL response received") + + var result struct { + Data map[string]any `json:"data"` + Errors []struct { + Message string `json:"message"` + Path []any `json:"path,omitempty"` + } `json:"errors,omitempty"` + } + + if err := json.Unmarshal(respBody, &result); err != nil { + return nil, fmt.Errorf("failed to parse GraphQL response: %w", err) + } + + if len(result.Errors) > 0 { + errMsgs := make([]string, len(result.Errors)) + for i, e := range result.Errors { + errMsgs[i] = e.Message + } + c.log.WithField("errors", errMsgs).Warn("GraphQL query returned errors") + if result.Data == nil { + return nil, fmt.Errorf("GraphQL errors: %v", errMsgs) + } + } + + return result.Data, nil +} + +// GetCurrentUser returns information about the current authenticated user. +func (c *InternalClient) GetCurrentUser(ctx context.Context) (*tools.UserInfo, error) { + c.log.Debug("getting current user from context") + + actor := authz.ActorFromContext(ctx) + if actor == nil || actor.User == nil { + return nil, fmt.Errorf("no authenticated user in context") + } + + return &tools.UserInfo{ + Name: actor.User.Identity(), + IsAdmin: actor.User.IsAdmin(), + }, nil +} + +// GetUserTeams returns the teams the current user belongs to. +func (c *InternalClient) GetUserTeams(ctx context.Context) ([]tools.TeamInfo, error) { + c.log.Debug("getting user teams") + + actor := authz.ActorFromContext(ctx) + if actor == nil || actor.User == nil { + return nil, fmt.Errorf("no authenticated user in context") + } + + userID := actor.User.GetID() + + first := 100 + page, err := pagination.ParsePage(&first, nil, nil, nil) + if err != nil { + return nil, fmt.Errorf("failed to parse pagination: %w", err) + } + + teamsResult, err := team.ListForUser(ctx, userID, page, nil) + if err != nil { + return nil, fmt.Errorf("failed to list teams for user: %w", err) + } + + nodes := teamsResult.Nodes() + result := make([]tools.TeamInfo, 0, len(nodes)) + for _, member := range nodes { + t, err := team.Get(ctx, member.TeamSlug) + if err != nil { + c.log.WithError(err).WithField("team_slug", member.TeamSlug).Warn("failed to get team details") + continue + } + + result = append(result, tools.TeamInfo{ + Slug: string(member.TeamSlug), + Purpose: t.Purpose, + Role: string(member.Role), + }) + } + + c.log.WithField("team_count", len(result)).Debug("retrieved user teams") + return result, nil +} diff --git a/internal/agent/internal_client_test.go b/internal/agent/internal_client_test.go new file mode 100644 index 000000000..6083772ed --- /dev/null +++ b/internal/agent/internal_client_test.go @@ -0,0 +1,104 @@ +package agent + +import ( + "context" + "testing" + + "github.com/google/uuid" + "github.com/nais/api/internal/auth/authz" + "github.com/nais/api/internal/slug" + "github.com/sirupsen/logrus" +) + +type mockAuthenticatedUser struct { + id uuid.UUID + identity string + isServiceAccount bool + isAdmin bool +} + +func (m *mockAuthenticatedUser) GetID() uuid.UUID { return m.id } +func (m *mockAuthenticatedUser) Identity() string { return m.identity } +func (m *mockAuthenticatedUser) IsServiceAccount() bool { return m.isServiceAccount } +func (m *mockAuthenticatedUser) IsAdmin() bool { return m.isAdmin } +func (m *mockAuthenticatedUser) GCPTeamGroups(_ context.Context) ([]string, error) { return nil, nil } + +func TestInternalClient_GetCurrentUser(t *testing.T) { + // Create a mock user + mockUser := &mockAuthenticatedUser{ + id: uuid.New(), + identity: "test@example.com", + isAdmin: false, + } + + // Create context with user + ctx := authz.ContextWithActor(context.Background(), mockUser, nil) + + // Create client (handler is nil for this test since we don't call ExecuteGraphQL) + client := &InternalClient{ + log: logrus.NewEntry(logrus.StandardLogger()), + } + + // Test GetCurrentUser + user, err := client.GetCurrentUser(ctx) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if user.Name != "test@example.com" { + t.Errorf("expected name %q, got %q", "test@example.com", user.Name) + } + + if user.IsAdmin { + t.Error("expected user to not be admin") + } +} + +func TestInternalClient_GetCurrentUser_NoUser(t *testing.T) { + // Create context without user + ctx := context.Background() + + // Create client + client := &InternalClient{ + log: logrus.NewEntry(logrus.StandardLogger()), + } + + // Test GetCurrentUser + _, err := client.GetCurrentUser(ctx) + if err == nil { + t.Error("expected error when no user in context") + } +} + +func TestInternalClient_GetCurrentUser_Admin(t *testing.T) { + // Create a mock admin user + mockUser := &mockAuthenticatedUser{ + id: uuid.New(), + identity: "admin@example.com", + isAdmin: true, + } + + // Create context with user + ctx := authz.ContextWithActor(context.Background(), mockUser, nil) + + // Create client + client := &InternalClient{ + log: logrus.NewEntry(logrus.StandardLogger()), + } + + // Test GetCurrentUser + user, err := client.GetCurrentUser(ctx) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if !user.IsAdmin { + t.Error("expected user to be admin") + } +} + +// Ensure mockAuthenticatedUser implements authz.AuthenticatedUser +var _ authz.AuthenticatedUser = (*mockAuthenticatedUser)(nil) + +// Dummy reference to slug package to avoid unused import error in tests +var _ = slug.Slug("") diff --git a/internal/agent/model.go b/internal/agent/model.go new file mode 100644 index 000000000..38f5d9321 --- /dev/null +++ b/internal/agent/model.go @@ -0,0 +1,47 @@ +package agent + +import ( + "time" + + "github.com/google/uuid" +) + +// ChatContext contains the user's current UI context. +type ChatContext struct { + Path string `json:"path,omitempty"` + Team string `json:"team,omitempty"` + App string `json:"app,omitempty"` + Env string `json:"env,omitempty"` +} + +// Source describes a documentation source used in the response. +type Source struct { + Title string `json:"title"` + URL string `json:"url"` +} + +// ConversationSummary represents a conversation in list view. +type ConversationSummary struct { + ID uuid.UUID `json:"id"` + Title string `json:"title"` + UpdatedAt time.Time `json:"updated_at"` +} + +// Conversation represents a full conversation with messages. +type Conversation struct { + ID uuid.UUID `json:"id"` + Title string `json:"title"` + Messages []ConversationMessage `json:"messages"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` +} + +// ConversationMessage represents a message in a conversation. +type ConversationMessage struct { + ID uuid.UUID `json:"id"` + Role string `json:"role"` + Content string `json:"content"` + Blocks []ContentBlock `json:"blocks,omitempty"` + Sources []Source `json:"sources,omitempty"` + CreatedAt time.Time `json:"created_at"` +} diff --git a/internal/agent/orchestrator.go b/internal/agent/orchestrator.go new file mode 100644 index 000000000..1832cc9f1 --- /dev/null +++ b/internal/agent/orchestrator.go @@ -0,0 +1,535 @@ +package agent + +import ( + "context" + "fmt" + "strings" + "time" + + "github.com/nais/api/internal/agent/chat" + "github.com/nais/api/internal/agent/rag" + "github.com/nais/api/internal/agent/tools" + "github.com/sirupsen/logrus" +) + +const ( + maxToolIterations = 5 + requestTimeout = 60 * time.Second + retryAttempts = 3 + retryBaseDelay = 100 * time.Millisecond + maxToolOutputChars = 50000 // Truncate tool outputs to prevent context window exhaustion +) + +// OrchestratorEventType is the type of a streaming event from the orchestrator. +// It uses the same string values as the SSE event type field, so no translation is needed. +type OrchestratorEventType = string + +const ( + OrchestratorEventToolStart OrchestratorEventType = "tool_start" + OrchestratorEventToolEnd OrchestratorEventType = "tool_end" + OrchestratorEventContent OrchestratorEventType = "content" + OrchestratorEventThinking OrchestratorEventType = "thinking" + OrchestratorEventChart OrchestratorEventType = "chart" + OrchestratorEventError OrchestratorEventType = "error" + OrchestratorEventDone OrchestratorEventType = "done" +) + +// ChartData represents the data needed to render a chart on the client. +type ChartData = tools.ChartData + +// ContentBlockType represents the type of content block in an assistant message. +type ContentBlockType string + +const ( + // ContentBlockTypeThinking represents the model's reasoning/thought process. + ContentBlockTypeThinking ContentBlockType = "thinking" + // ContentBlockTypeText represents regular text output. + ContentBlockTypeText ContentBlockType = "text" + // ContentBlockTypeToolUse represents a tool invocation and its result. + // Stored internally for LLM history reconstruction; filtered from client responses. + ContentBlockTypeToolUse ContentBlockType = "tool_use" + // ContentBlockTypeChart represents a chart to be rendered. + ContentBlockTypeChart ContentBlockType = "chart" +) + +// ContentBlock represents a single block of content in an assistant message. +// Messages are composed of multiple blocks displayed in order. +// +// Blocks of type "text", "thinking", and "chart" are shown to the client. +// Blocks of type "tool_use" are stored for LLM history reconstruction only +// and should be filtered out of client-visible responses. +type ContentBlock struct { + // Type indicates the kind of content block. + Type ContentBlockType `json:"type"` + // Text is the text content (for "text" type blocks). + Text string `json:"text,omitempty"` + // Thinking is the model's reasoning (for "thinking" type blocks). + Thinking string `json:"thinking,omitempty"` + // ToolCallID is the unique identifier for the tool call (for "tool_use" type blocks). + ToolCallID string `json:"tool_call_id,omitempty"` + // ToolName is the name of the tool (for "tool_use" type blocks). + ToolName string `json:"tool_name,omitempty"` + // ToolArguments are the arguments passed to the tool (for "tool_use" type blocks). + // Stored so the full tool call can be reconstructed when replaying history to the LLM. + ToolArguments map[string]any `json:"tool_arguments,omitempty"` + // ToolSuccess indicates whether the tool execution succeeded (for "tool_use" type blocks). + ToolSuccess bool `json:"tool_success,omitempty"` + // ToolResult is the result returned by the tool (for "tool_use" type blocks). + // Stored to reconstruct the full conversation history for subsequent LLM calls. + ToolResult string `json:"tool_result,omitempty"` + // Chart contains chart data (for "chart" type blocks). + Chart *ChartData `json:"chart,omitempty"` +} + +// OrchestratorStreamEvent represents an event in the streaming response from the orchestrator. +// Type matches the SSE wire format strings directly. +type OrchestratorStreamEvent struct { + Type OrchestratorEventType + Content string + Thinking string + ToolName string + Description string + Success bool + Error error + Chart *ChartData + Usage *chat.UsageStats + ToolCallID string + // Result is only set for OrchestratorEventDone and carries the complete + // accumulated result (blocks + usage) for persistence. + Result *OrchestratorResult +} + +// ToStreamEvent converts an OrchestratorStreamEvent to a StreamEvent for SSE delivery. +// It handles all regular event types via direct field mapping. The caller is responsible +// for handling OrchestratorEventError and OrchestratorEventDone, which require control +// flow (return) and additional side effects (logging, storing messages, etc.). +func (e OrchestratorStreamEvent) ToStreamEvent() StreamEvent { + return StreamEvent{ + Type: e.Type, + Content: e.Content, + Thinking: e.Thinking, + ToolName: e.ToolName, + ToolCallID: e.ToolCallID, + Description: e.Description, + ToolSuccess: e.Success, + Chart: e.Chart, + Usage: e.Usage, + } +} + +// OrchestratorResult contains the complete result of an orchestrated conversation turn. +type OrchestratorResult struct { + // Blocks contains the sequence of content blocks that make up the response. + // Includes "text", "thinking", "chart", and internal "tool_use" blocks. + Blocks []ContentBlock + // Usage contains aggregated token usage statistics. + Usage *chat.UsageStats +} + +// Orchestrator manages the conversation loop between user, LLM, and tools. +type Orchestrator struct { + chatClient chat.StreamingClient + registry *tools.Registry + log logrus.FieldLogger +} + +// NewOrchestrator creates a new orchestrator. +func NewOrchestrator( + chatClient chat.StreamingClient, + registry *tools.Registry, + log logrus.FieldLogger, +) *Orchestrator { + return &Orchestrator{ + chatClient: chatClient, + registry: registry, + log: log, + } +} + +// conversationLoop manages the shared mutable state for a single conversation turn, +// used by both Run and RunStream. +type conversationLoop struct { + orchestrator *Orchestrator + ctx context.Context + messages []chat.Message + tools []chat.ToolDefinition + systemPrompt string + docs []rag.Document + blocks []ContentBlock + totalUsage chat.UsageStats + log logrus.FieldLogger +} + +func (o *Orchestrator) newConversationLoop( + ctx context.Context, + userMessage string, + chatCtx *ChatContext, + docs []rag.Document, + history []chat.Message, +) *conversationLoop { + toolDefs := o.getToolDefinitions() + systemPrompt := o.buildSystemPrompt(chatCtx, docs, toolDefs) + + messages := make([]chat.Message, 0, len(history)+1) + messages = append(messages, history...) + messages = append(messages, chat.Message{ + Role: chat.RoleUser, + Content: userMessage, + }) + + return &conversationLoop{ + orchestrator: o, + ctx: ctx, + messages: messages, + tools: toolDefs, + systemPrompt: systemPrompt, + docs: docs, + log: o.log, + } +} + +func (cl *conversationLoop) buildRequest() *chat.Request { + return &chat.Request{ + SystemPrompt: cl.systemPrompt, + Messages: cl.messages, + Tools: cl.tools, + Documents: cl.docs, + } +} + +func (cl *conversationLoop) addThinkingBlock(thinking string) { + if thinking != "" { + cl.blocks = append(cl.blocks, ContentBlock{ + Type: ContentBlockTypeThinking, + Thinking: thinking, + }) + } +} + +func (cl *conversationLoop) addTextBlock(text string) { + if text != "" { + cl.blocks = append(cl.blocks, ContentBlock{ + Type: ContentBlockTypeText, + Text: text, + }) + } +} + +func (cl *conversationLoop) addAssistantMessage(content string, toolCalls []chat.ToolCall) { + cl.messages = append(cl.messages, chat.Message{ + Role: chat.RoleAssistant, + Content: content, + ToolCalls: toolCalls, + }) +} + +// toolExecutionResult holds the result of executing a single tool call. +type toolExecutionResult struct { + ToolCall chat.ToolCall + Result string + ChartData *ChartData + Success bool + Error error +} + +func (cl *conversationLoop) executeToolCall(toolCall chat.ToolCall) toolExecutionResult { + cl.log.WithFields(logrus.Fields{ + "tool": toolCall.Name, + "tool_id": toolCall.ID, + }).Debug("executing tool call") + + toolResult, chartData, err := cl.orchestrator.executeTool(cl.ctx, toolCall) + success := err == nil + + var resultContent string + if err != nil { + resultContent = fmt.Sprintf("Error executing tool: %s", err.Error()) + cl.log.WithError(err).WithField("tool", toolCall.Name).Warn("tool execution failed") + } else { + cl.log.WithFields(logrus.Fields{ + "tool": toolCall.Name, + "result_length": len(toolResult), + "has_chart": chartData != nil, + }).Debug("tool execution succeeded") + resultContent = toolResult + } + + return toolExecutionResult{ + ToolCall: toolCall, + Result: resultContent, + ChartData: chartData, + Success: success, + Error: err, + } +} + +// recordToolExecution stores a tool execution in the block history and appends the +// tool result message for the next LLM iteration. +func (cl *conversationLoop) recordToolExecution(result toolExecutionResult) { + cl.blocks = append(cl.blocks, ContentBlock{ + Type: ContentBlockTypeToolUse, + ToolCallID: result.ToolCall.ID, + ToolName: result.ToolCall.Name, + ToolArguments: result.ToolCall.Arguments, + ToolSuccess: result.Success, + ToolResult: result.Result, + }) + + if result.ChartData != nil { + cl.blocks = append(cl.blocks, ContentBlock{ + Type: ContentBlockTypeChart, + Chart: result.ChartData, + }) + } + + cl.messages = append(cl.messages, chat.Message{ + Role: chat.RoleTool, + Content: result.Result, + ToolCallID: result.ToolCall.ID, + }) +} + +func (cl *conversationLoop) accumulateUsage(usage *chat.UsageStats) { + if usage != nil { + cl.totalUsage.InputTokens += usage.InputTokens + cl.totalUsage.OutputTokens += usage.OutputTokens + cl.totalUsage.TotalTokens += usage.TotalTokens + if usage.MaxTokens > 0 { + cl.totalUsage.MaxTokens = usage.MaxTokens + } + } +} + +func (cl *conversationLoop) result() *OrchestratorResult { + return &OrchestratorResult{ + Blocks: cl.blocks, + Usage: &cl.totalUsage, + } +} + +// Run executes a non-streaming conversation turn and returns the complete result. +func (o *Orchestrator) Run( + ctx context.Context, + userMessage string, + chatCtx *ChatContext, + docs []rag.Document, + history []chat.Message, +) (*OrchestratorResult, error) { + ctx, cancel := context.WithTimeout(ctx, requestTimeout) + defer cancel() + + o.log.WithFields(logrus.Fields{ + "message_length": len(userMessage), + "doc_count": len(docs), + "history_length": len(history), + }).Debug("starting orchestrator run") + + loop := o.newConversationLoop(ctx, userMessage, chatCtx, docs, history) + + for iteration := 0; iteration < maxToolIterations; iteration++ { + o.log.WithField("iteration", iteration).Debug("starting tool iteration") + + resp, err := o.callLLMWithRetry(ctx, loop.buildRequest()) + if err != nil { + return nil, fmt.Errorf("LLM call failed: %w", err) + } + + loop.accumulateUsage(resp.Usage) + + if len(resp.ToolCalls) == 0 { + loop.addThinkingBlock(resp.Thinking) + loop.addTextBlock(resp.Content) + return loop.result(), nil + } + + loop.addThinkingBlock(resp.Thinking) + loop.addTextBlock(resp.Content) + loop.addAssistantMessage(resp.Content, resp.ToolCalls) + + for _, toolCall := range resp.ToolCalls { + result := loop.executeToolCall(toolCall) + loop.recordToolExecution(result) + } + } + + return nil, fmt.Errorf("max tool iterations (%d) exceeded", maxToolIterations) +} + +// RunStream executes a streaming conversation turn. Events are sent to the returned channel. +// The final event is always OrchestratorEventDone, which carries the complete OrchestratorResult +// for persistence. The channel is closed after the Done event is sent. +func (o *Orchestrator) RunStream( + ctx context.Context, + userMessage string, + chatCtx *ChatContext, + docs []rag.Document, + history []chat.Message, +) <-chan OrchestratorStreamEvent { + eventCh := make(chan OrchestratorStreamEvent, 100) + + go func() { + defer close(eventCh) + + ctx, cancel := context.WithTimeout(ctx, requestTimeout) + defer cancel() + + loop := o.newConversationLoop(ctx, userMessage, chatCtx, docs, history) + + for iteration := 0; iteration < maxToolIterations; iteration++ { + o.log.WithField("iteration", iteration).Debug("starting streaming tool iteration") + + streamCh, err := o.chatClient.ChatStream(ctx, loop.buildRequest()) + if err != nil { + eventCh <- OrchestratorStreamEvent{ + Type: OrchestratorEventError, + Error: fmt.Errorf("LLM stream failed: %w", err), + } + return + } + + content, thinking, toolCalls, usage, abort := o.processStream(streamCh, eventCh) + if abort { + return + } + + loop.accumulateUsage(usage) + + if len(toolCalls) == 0 { + // No tool calls — this is the final response. + loop.addThinkingBlock(thinking) + loop.addTextBlock(content) + result := loop.result() + eventCh <- OrchestratorStreamEvent{ + Type: OrchestratorEventDone, + Usage: result.Usage, + Result: result, + } + return + } + + // Flush accumulated content/thinking to blocks before executing tools. + loop.addThinkingBlock(thinking) + loop.addTextBlock(content) + loop.addAssistantMessage(content, toolCalls) + + o.executeToolCallsStreaming(loop, toolCalls, eventCh) + } + + eventCh <- OrchestratorStreamEvent{ + Type: OrchestratorEventError, + Error: fmt.Errorf("max tool iterations (%d) exceeded", maxToolIterations), + } + }() + + return eventCh +} + +// processStream consumes the LLM stream, forwarding events and accumulating results. +// Returns accumulated content, thinking, tool calls, usage, and whether to abort (due to error). +func (o *Orchestrator) processStream( + streamCh <-chan chat.StreamChunk, + eventCh chan<- OrchestratorStreamEvent, +) (content string, thinking string, toolCalls []chat.ToolCall, usage *chat.UsageStats, abort bool) { + var contentBuilder strings.Builder + var thinkingBuilder strings.Builder + usage = &chat.UsageStats{} + + for chunk := range streamCh { + if chunk.Error != nil { + o.log.WithError(chunk.Error).Error("received error in stream chunk") + eventCh <- OrchestratorStreamEvent{ + Type: OrchestratorEventError, + Error: chunk.Error, + } + return "", "", nil, nil, true + } + + if chunk.Usage != nil { + usage.InputTokens += chunk.Usage.InputTokens + usage.OutputTokens += chunk.Usage.OutputTokens + usage.TotalTokens += chunk.Usage.TotalTokens + if chunk.Usage.MaxTokens > 0 { + usage.MaxTokens = chunk.Usage.MaxTokens + } + } + + if chunk.Thinking != "" { + thinkingBuilder.WriteString(chunk.Thinking) + eventCh <- OrchestratorStreamEvent{ + Type: OrchestratorEventThinking, + Thinking: chunk.Thinking, + } + } + + if chunk.Content != "" { + contentBuilder.WriteString(chunk.Content) + eventCh <- OrchestratorStreamEvent{ + Type: OrchestratorEventContent, + Content: chunk.Content, + } + } + + if len(chunk.ToolCalls) > 0 { + toolCalls = append(toolCalls, chunk.ToolCalls...) + } + } + + return contentBuilder.String(), thinkingBuilder.String(), toolCalls, usage, false +} + +// executeToolCallsStreaming executes tool calls, emits streaming events, and records results +// in the conversation loop for the next LLM iteration. +func (o *Orchestrator) executeToolCallsStreaming( + loop *conversationLoop, + toolCalls []chat.ToolCall, + eventCh chan<- OrchestratorStreamEvent, +) { + for _, toolCall := range toolCalls { + eventCh <- OrchestratorStreamEvent{ + Type: OrchestratorEventToolStart, + ToolName: toolCall.Name, + ToolCallID: toolCall.ID, + Description: fmt.Sprintf("Executing %s...", toolCall.Name), + } + + result := loop.executeToolCall(toolCall) + + eventCh <- OrchestratorStreamEvent{ + Type: OrchestratorEventToolEnd, + ToolName: toolCall.Name, + ToolCallID: toolCall.ID, + Description: fmt.Sprintf("Executed %s", toolCall.Name), + Success: result.Success, + } + + if result.ChartData != nil { + eventCh <- OrchestratorStreamEvent{ + Type: OrchestratorEventChart, + Chart: result.ChartData, + } + } + + // Record in the loop (adds tool_use block + tool result message for next iteration). + loop.recordToolExecution(result) + } +} + +func (o *Orchestrator) callLLMWithRetry(ctx context.Context, req *chat.Request) (*chat.Response, error) { + var lastErr error + for attempt := 0; attempt < retryAttempts; attempt++ { + resp, err := o.chatClient.Chat(ctx, req) + if err == nil { + return resp, nil + } + lastErr = err + o.log.WithError(err).WithField("attempt", attempt+1).Warn("LLM call failed, retrying") + + delay := retryBaseDelay * time.Duration(1< 0 { + sb.WriteString("\n## Documentation\n") + for _, doc := range docs { + sb.WriteString(fmt.Sprintf("\n### %s\n%s\nSource: %s\n", doc.Title, doc.Content, doc.URL)) + } + } + + return sb.String() +} + +// firstSentence extracts the first sentence from a string for concise tool descriptions. +func firstSentence(s string) string { + // Find the first period followed by a space or end of string + for i := 0; i < len(s)-1; i++ { + if s[i] == '.' && (i+1 >= len(s) || s[i+1] == ' ' || s[i+1] == '\n') { + return s[:i+1] + } + } + // If no sentence ending found, return the whole string (truncated if too long) + if len(s) > 150 { + return s[:147] + "..." + } + return s +} diff --git a/internal/agent/orchestrator_test.go b/internal/agent/orchestrator_test.go new file mode 100644 index 000000000..dae5a96b7 --- /dev/null +++ b/internal/agent/orchestrator_test.go @@ -0,0 +1,135 @@ +package agent + +import ( + "testing" + + "github.com/nais/api/internal/agent/chat" +) + +func TestExtractToolCallsFromBlocks(t *testing.T) { + tests := []struct { + name string + blocks []ContentBlock + want []chat.ToolCall + }{ + { + name: "empty blocks", + blocks: []ContentBlock{}, + want: nil, + }, + { + name: "no tool blocks", + blocks: []ContentBlock{ + {Type: ContentBlockTypeText, Text: "Hello"}, + {Type: ContentBlockTypeThinking, Thinking: "Thinking..."}, + }, + want: nil, + }, + { + name: "single tool call", + blocks: []ContentBlock{ + {Type: ContentBlockTypeToolUse, ToolCallID: "call_1", ToolName: "execute_graphql", ToolSuccess: true, ToolResult: "result data"}, + }, + want: []chat.ToolCall{ + {ID: "call_1", Name: "execute_graphql"}, + }, + }, + { + name: "multiple tool calls in order", + blocks: []ContentBlock{ + {Type: ContentBlockTypeText, Text: "Let me check..."}, + {Type: ContentBlockTypeToolUse, ToolCallID: "call_1", ToolName: "schema_get_type", ToolSuccess: true, ToolResult: "type info"}, + {Type: ContentBlockTypeToolUse, ToolCallID: "call_2", ToolName: "execute_graphql", ToolSuccess: true, ToolResult: "query result"}, + {Type: ContentBlockTypeText, Text: "Here's what I found..."}, + }, + want: []chat.ToolCall{ + {ID: "call_1", Name: "schema_get_type"}, + {ID: "call_2", Name: "execute_graphql"}, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := extractToolCallsFromBlocks(tt.blocks) + if len(got) != len(tt.want) { + t.Errorf("extractToolCallsFromBlocks() returned %d tool calls, want %d", len(got), len(tt.want)) + return + } + for i, tc := range got { + if tc.ID != tt.want[i].ID || tc.Name != tt.want[i].Name { + t.Errorf("extractToolCallsFromBlocks()[%d] = {ID: %q, Name: %q}, want {ID: %q, Name: %q}", + i, tc.ID, tc.Name, tt.want[i].ID, tt.want[i].Name) + } + } + }) + } +} + +func TestExtractTextContentFromBlocks(t *testing.T) { + tests := []struct { + name string + blocks []ContentBlock + want string + }{ + { + name: "empty blocks", + blocks: []ContentBlock{}, + want: "", + }, + { + name: "single text block", + blocks: []ContentBlock{ + {Type: ContentBlockTypeText, Text: "Hello, world!"}, + }, + want: "Hello, world!", + }, + { + name: "multiple text blocks", + blocks: []ContentBlock{ + {Type: ContentBlockTypeText, Text: "First part."}, + {Type: ContentBlockTypeToolUse, ToolName: "some_tool"}, + {Type: ContentBlockTypeText, Text: "Second part."}, + }, + want: "First part. Second part.", + }, + { + name: "ignores non-text blocks", + blocks: []ContentBlock{ + {Type: ContentBlockTypeThinking, Thinking: "Thinking..."}, + {Type: ContentBlockTypeToolUse, ToolName: "tool", ToolResult: "result"}, + {Type: ContentBlockTypeChart, Chart: &ChartData{Title: "Chart"}}, + }, + want: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := extractTextContentFromBlocks(tt.blocks) + if got != tt.want { + t.Errorf("extractTextContentFromBlocks() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestClientVisibleBlocks(t *testing.T) { + blocks := []ContentBlock{ + {Type: ContentBlockTypeThinking, Thinking: "reasoning"}, + {Type: ContentBlockTypeText, Text: "here is the answer"}, + {Type: ContentBlockTypeToolUse, ToolCallID: "call_1", ToolName: "execute_graphql", ToolSuccess: true, ToolResult: `{"data":{}}`}, + {Type: ContentBlockTypeChart, Chart: &ChartData{Title: "CPU Usage"}}, + } + + got := clientVisibleBlocks(blocks) + + if len(got) != 3 { + t.Fatalf("expected 3 visible blocks, got %d", len(got)) + } + for _, b := range got { + if b.Type == ContentBlockTypeToolUse { + t.Errorf("clientVisibleBlocks() should not include tool_use blocks") + } + } +} diff --git a/internal/agent/orchestrator_tools.go b/internal/agent/orchestrator_tools.go new file mode 100644 index 000000000..b78f6006c --- /dev/null +++ b/internal/agent/orchestrator_tools.go @@ -0,0 +1,75 @@ +package agent + +import ( + "context" + "encoding/json" + "fmt" + + "github.com/nais/api/internal/agent/chat" + "github.com/sirupsen/logrus" +) + +// executeTool executes a tool call using the registry. +// Returns the result string and an optional ChartData if this was a render_chart call. +func (o *Orchestrator) executeTool(ctx context.Context, toolCall chat.ToolCall) (string, *ChartData, error) { + o.log.WithFields(logrus.Fields{ + "tool": toolCall.Name, + "args": toolCall.Arguments, + }).Debug("executing tool") + + result, err := o.registry.Execute(ctx, toolCall.Name, toolCall.Arguments) + if err != nil { + return "", nil, fmt.Errorf("tool %s failed: %w", toolCall.Name, err) + } + + // Check if the result is ChartData (handled specially for visualization) + if chart, ok := result.(*ChartData); ok { + return "Chart rendered successfully. The user can now see the visualization.", chart, nil + } + + // Convert result to JSON string for LLM consumption + resultJSON, err := json.Marshal(result) + if err != nil { + return "", nil, fmt.Errorf("failed to marshal tool result: %w", err) + } + + output := string(resultJSON) + + // Truncate large outputs to prevent context window exhaustion + if len(output) > maxToolOutputChars { + o.log.WithFields(logrus.Fields{ + "tool": toolCall.Name, + "original_length": len(output), + "truncated_to": maxToolOutputChars, + }).Warn("truncating large tool output") + output = output[:maxToolOutputChars] + "\n\n[Output truncated due to size. Please refine your query to get more specific results.]" + } + + return output, nil, nil +} + +// getToolDefinitions returns the tool definitions for the LLM. +func (o *Orchestrator) getToolDefinitions() []chat.ToolDefinition { + registeredTools := o.registry.ListTools() + result := make([]chat.ToolDefinition, 0, len(registeredTools)) + + for _, tool := range registeredTools { + chatParams := make([]chat.ParameterDefinition, len(tool.Parameters)) + for i, p := range tool.Parameters { + chatParams[i] = chat.ParameterDefinition{ + Name: p.Name, + Type: p.Type, + Description: p.Description, + Required: p.Required, + } + } + + result = append(result, chat.ToolDefinition{ + Name: tool.Name, + Description: tool.Description, + Parameters: chatParams, + }) + } + + return result +} diff --git a/internal/agent/queries.go b/internal/agent/queries.go new file mode 100644 index 000000000..16de2ced4 --- /dev/null +++ b/internal/agent/queries.go @@ -0,0 +1,339 @@ +package agent + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5" + "github.com/nais/api/internal/agent/agentsql" + "github.com/nais/api/internal/agent/chat" + "github.com/nais/api/internal/database" +) + +// ErrConversationNotFound is returned when a conversation does not exist or does not belong to the user. +var ErrConversationNotFound = errors.New("conversation not found") + +// GetOrCreateConversation retrieves an existing conversation or creates a new one. +func GetOrCreateConversation(ctx context.Context, userID uuid.UUID, conversationID string, firstMessage string) (uuid.UUID, error) { + if conversationID != "" { + id, err := uuid.Parse(conversationID) + if err != nil { + return uuid.Nil, fmt.Errorf("invalid conversation ID: %w", err) + } + + exists, err := db(ctx).ConversationExists(ctx, agentsql.ConversationExistsParams{ + ID: id, + UserID: userID, + }) + if err != nil { + return uuid.Nil, fmt.Errorf("failed to check conversation: %w", err) + } + if !exists { + return uuid.Nil, ErrConversationNotFound + } + + if err := db(ctx).TouchConversation(ctx, id); err != nil { + return uuid.Nil, fmt.Errorf("failed to update conversation: %w", err) + } + + return id, nil + } + + return createConversation(ctx, userID, firstMessage) +} + +func createConversation(ctx context.Context, userID uuid.UUID, firstMessage string) (uuid.UUID, error) { + id, err := db(ctx).CreateConversation(ctx, agentsql.CreateConversationParams{ + UserID: userID, + Title: generateTitle(firstMessage), + }) + if err != nil { + return uuid.Nil, fmt.Errorf("failed to create conversation: %w", err) + } + return id, nil +} + +// StoreMessages stores the user message and assistant response in the conversation. +// sources are the RAG documents surfaced during this turn and are stored with the +// assistant message so they can be returned when the conversation is loaded later. +func StoreMessages(ctx context.Context, conversationID uuid.UUID, userMessage string, result *OrchestratorResult, sources []Source) error { + return database.Transaction(ctx, func(ctx context.Context) error { + if err := db(ctx).InsertMessage(ctx, agentsql.InsertMessageParams{ + ConversationID: conversationID, + Role: "user", + Content: userMessage, + Blocks: nil, + Sources: nil, + }); err != nil { + return fmt.Errorf("failed to store user message: %w", err) + } + + var blocksJSON []byte + if len(result.Blocks) > 0 { + var err error + blocksJSON, err = json.Marshal(result.Blocks) + if err != nil { + return fmt.Errorf("failed to marshal blocks: %w", err) + } + } + + var sourcesJSON []byte + if len(sources) > 0 { + var err error + sourcesJSON, err = json.Marshal(sources) + if err != nil { + return fmt.Errorf("failed to marshal sources: %w", err) + } + } + + if err := db(ctx).InsertMessage(ctx, agentsql.InsertMessageParams{ + ConversationID: conversationID, + Role: "assistant", + Content: "", + Blocks: blocksJSON, + Sources: sourcesJSON, + }); err != nil { + return fmt.Errorf("failed to store assistant message: %w", err) + } + + if err := db(ctx).TouchConversation(ctx, conversationID); err != nil { + return fmt.Errorf("failed to update conversation timestamp: %w", err) + } + + return nil + }) +} + +// ListConversations returns all conversations for a user, ordered by most recent. +func ListConversations(ctx context.Context, userID uuid.UUID) ([]ConversationSummary, error) { + rows, err := db(ctx).ListConversations(ctx, userID) + if err != nil { + return nil, fmt.Errorf("failed to query conversations: %w", err) + } + + conversations := make([]ConversationSummary, 0, len(rows)) + for _, row := range rows { + conversations = append(conversations, ConversationSummary{ + ID: row.ID, + Title: row.Title, + UpdatedAt: row.UpdatedAt.Time, + }) + } + + return conversations, nil +} + +// GetConversation retrieves a full conversation with all messages. +// Tool-use blocks are filtered from the returned messages since they are +// internal LLM plumbing and not meaningful to clients. +func GetConversation(ctx context.Context, userID uuid.UUID, conversationID uuid.UUID) (*Conversation, error) { + row, err := db(ctx).GetConversation(ctx, agentsql.GetConversationParams{ + ID: conversationID, + UserID: userID, + }) + if err != nil { + if err == pgx.ErrNoRows { + return nil, ErrConversationNotFound + } + return nil, fmt.Errorf("failed to get conversation: %w", err) + } + + conv := &Conversation{ + ID: row.ID, + Title: row.Title, + CreatedAt: row.CreatedAt.Time, + UpdatedAt: row.UpdatedAt.Time, + } + + messages, err := db(ctx).GetConversationMessages(ctx, conversationID) + if err != nil { + return nil, fmt.Errorf("failed to query messages: %w", err) + } + + conv.Messages = make([]ConversationMessage, 0, len(messages)) + for _, msg := range messages { + cm := ConversationMessage{ + ID: msg.ID, + Role: msg.Role, + Content: msg.Content, + CreatedAt: msg.CreatedAt.Time, + } + + if msg.Blocks != nil { + var allBlocks []ContentBlock + if err := json.Unmarshal(msg.Blocks, &allBlocks); err == nil { + // Filter out tool_use blocks — they are internal LLM history details. + cm.Blocks = clientVisibleBlocks(allBlocks) + + if msg.Role == "assistant" && cm.Content == "" { + cm.Content = extractTextContentFromBlocks(allBlocks) + } + } + } + + if msg.Sources != nil { + var srcs []Source + if err := json.Unmarshal(msg.Sources, &srcs); err == nil { + cm.Sources = srcs + } + } + + conv.Messages = append(conv.Messages, cm) + } + + return conv, nil +} + +// DeleteConversation deletes a conversation and all its messages. +func DeleteConversation(ctx context.Context, userID uuid.UUID, conversationID uuid.UUID) error { + rowsAffected, err := db(ctx).DeleteConversation(ctx, agentsql.DeleteConversationParams{ + ID: conversationID, + UserID: userID, + }) + if err != nil { + return fmt.Errorf("failed to delete conversation: %w", err) + } + + if rowsAffected == 0 { + return ErrConversationNotFound + } + + return nil +} + +// GetConversationHistory retrieves the message history for a conversation, formatted for the LLM. +// This reconstructs the full conversation including tool call responses from stored blocks. +func GetConversationHistory(ctx context.Context, conversationID uuid.UUID) ([]chat.Message, error) { + rows, err := db(ctx).GetConversationHistory(ctx, conversationID) + if err != nil { + return nil, fmt.Errorf("failed to query messages: %w", err) + } + + messages := make([]chat.Message, 0, len(rows)*2) + for _, row := range rows { + if row.Role == "user" { + messages = append(messages, chat.Message{ + Role: chat.RoleUser, + Content: row.Content, + }) + continue + } + + if row.Role == "assistant" && row.Blocks != nil { + var blocks []ContentBlock + if err := json.Unmarshal(row.Blocks, &blocks); err != nil { + messages = append(messages, chat.Message{ + Role: chat.RoleAssistant, + Content: row.Content, + }) + continue + } + + toolCalls := extractToolCallsFromBlocks(blocks) + textContent := extractTextContentFromBlocks(blocks) + + if len(toolCalls) > 0 { + messages = append(messages, chat.Message{ + Role: chat.RoleAssistant, + Content: "", + ToolCalls: toolCalls, + }) + + for _, block := range blocks { + if block.Type == ContentBlockTypeToolUse && block.ToolResult != "" { + messages = append(messages, chat.Message{ + Role: chat.RoleTool, + Content: block.ToolResult, + ToolCallID: block.ToolCallID, + }) + } + } + + if textContent != "" { + messages = append(messages, chat.Message{ + Role: chat.RoleAssistant, + Content: textContent, + }) + } + } else { + messages = append(messages, chat.Message{ + Role: chat.RoleAssistant, + Content: textContent, + }) + } + continue + } + + messages = append(messages, chat.Message{ + Role: chat.Role(row.Role), + Content: row.Content, + }) + } + + return messages, nil +} + +// clientVisibleBlocks returns only the blocks that are meaningful to display to a client, +// filtering out internal tool_use blocks. +func clientVisibleBlocks(blocks []ContentBlock) []ContentBlock { + result := make([]ContentBlock, 0, len(blocks)) + for _, b := range blocks { + if b.Type != ContentBlockTypeToolUse { + result = append(result, b) + } + } + return result +} + +// extractToolCallsFromBlocks extracts tool call information from blocks for LLM history. +func extractToolCallsFromBlocks(blocks []ContentBlock) []chat.ToolCall { + var toolCalls []chat.ToolCall + for _, block := range blocks { + if block.Type == ContentBlockTypeToolUse { + toolCalls = append(toolCalls, chat.ToolCall{ + ID: block.ToolCallID, + Name: block.ToolName, + Arguments: block.ToolArguments, + }) + } + } + return toolCalls +} + +// extractTextContentFromBlocks extracts all text content from blocks. +func extractTextContentFromBlocks(blocks []ContentBlock) string { + var sb strings.Builder + for _, block := range blocks { + if block.Type == ContentBlockTypeText && block.Text != "" { + if sb.Len() > 0 { + sb.WriteString(" ") + } + sb.WriteString(block.Text) + } + } + return sb.String() +} + +// generateTitle creates a title from the first message. +func generateTitle(message string) string { + title := strings.TrimSpace(message) + + if idx := strings.Index(title, "\n"); idx > 0 { + title = title[:idx] + } + + const maxLen = 100 + if len(title) > maxLen { + title = title[:maxLen-3] + "..." + } + + if title == "" { + title = "New conversation" + } + + return title +} diff --git a/internal/agent/queries/agent.sql b/internal/agent/queries/agent.sql new file mode 100644 index 000000000..d6ae403c9 --- /dev/null +++ b/internal/agent/queries/agent.sql @@ -0,0 +1,98 @@ +-- name: ConversationExists :one +SELECT + EXISTS ( + SELECT + 1 + FROM + agent_conversations + WHERE + id = @id + AND user_id = @user_id + ) +; + +-- name: TouchConversation :exec +UPDATE agent_conversations +SET + updated_at = NOW() +WHERE + id = @id +; + +-- name: CreateConversation :one +INSERT INTO + agent_conversations (user_id, title) +VALUES + (@user_id, @title) +RETURNING + id +; + +-- name: InsertMessage :exec +INSERT INTO + agent_messages (conversation_id, role, content, sources, blocks) +VALUES + (@conversation_id, @role, @content, @sources, @blocks) +; + +-- name: ListConversations :many +SELECT + id, + title, + updated_at +FROM + agent_conversations +WHERE + user_id = @user_id +ORDER BY + updated_at DESC +; + +-- name: GetConversation :one +SELECT + id, + title, + created_at, + updated_at +FROM + agent_conversations +WHERE + id = @id + AND user_id = @user_id +; + +-- name: GetConversationMessages :many +SELECT + id, + role, + content, + sources, + blocks, + created_at +FROM + agent_messages +WHERE + conversation_id = @conversation_id +ORDER BY + created_at ASC +; + +-- name: DeleteConversation :execrows +DELETE FROM agent_conversations +WHERE + id = @id + AND user_id = @user_id +; + +-- name: GetConversationHistory :many +SELECT + role, + content, + blocks +FROM + agent_messages +WHERE + conversation_id = @conversation_id +ORDER BY + created_at ASC +; diff --git a/internal/agent/rag/duckdb/searcher.go b/internal/agent/rag/duckdb/searcher.go new file mode 100644 index 000000000..6c33538bd --- /dev/null +++ b/internal/agent/rag/duckdb/searcher.go @@ -0,0 +1,246 @@ +// Package duckdb implements the rag.DocumentSearcher interface using a local DuckDB file. +package duckdb + +import ( + "context" + "database/sql" + "fmt" + "math" + "sort" + + _ "github.com/duckdb/duckdb-go/v2" + "github.com/nais/api/internal/agent/rag" + "github.com/sirupsen/logrus" + "google.golang.org/genai" +) + +// Config holds configuration for the DuckDB searcher. +type Config struct { + // DBPath is the path to the DuckDB file. + DBPath string + + // ProjectID is the GCP project ID for Vertex AI embeddings. + ProjectID string + + // Location is the GCP region for Vertex AI (must be in EU, e.g., "europe-west1"). + Location string + + // EmbeddingModel is the model to use for embeddings (e.g., "gemini-embedding-001"). + EmbeddingModel string +} + +// Searcher implements rag.DocumentSearcher using a local DuckDB file. +type Searcher struct { + db *sql.DB + embeddingModel string + genaiClient *genai.Client + log logrus.FieldLogger +} + +// NewSearcher creates a new DuckDB-based searcher. +func NewSearcher(ctx context.Context, cfg Config, log logrus.FieldLogger) (*Searcher, error) { + // Open DuckDB in read-only mode + db, err := sql.Open("duckdb", cfg.DBPath+"?access_mode=read_only") + if err != nil { + return nil, fmt.Errorf("failed to open DuckDB: %w", err) + } + + // Verify connection + if err := db.Ping(); err != nil { + db.Close() + return nil, fmt.Errorf("failed to ping DuckDB: %w", err) + } + + // Create Vertex AI client for embeddings + genaiClient, err := genai.NewClient(ctx, &genai.ClientConfig{ + Project: cfg.ProjectID, + Location: cfg.Location, + Backend: genai.BackendVertexAI, + }) + if err != nil { + db.Close() + return nil, fmt.Errorf("failed to create Vertex AI client: %w", err) + } + + log.WithFields(logrus.Fields{ + "db_path": cfg.DBPath, + "project": cfg.ProjectID, + "location": cfg.Location, + "embedding_model": cfg.EmbeddingModel, + }).Info("initialized DuckDB searcher with Vertex AI embeddings") + + return &Searcher{ + db: db, + embeddingModel: cfg.EmbeddingModel, + genaiClient: genaiClient, + log: log, + }, nil +} + +// Search queries the document store and returns relevant excerpts. +func (s *Searcher) Search(ctx context.Context, query string, opts *rag.SearchOptions) (*rag.SearchResult, error) { + maxResults := 5 + if opts != nil && opts.MaxResults > 0 { + maxResults = opts.MaxResults + } + + // Embed the query + queryEmbedding, err := s.embed(ctx, query) + if err != nil { + return nil, fmt.Errorf("failed to embed query: %w", err) + } + + // Fetch all documents with embeddings + // For small indices (~15MB), loading all into memory is acceptable + // Future optimization: use DuckDB vector extensions for ANN search + rows, err := s.db.QueryContext(ctx, ` + SELECT title, url, content, embedding + FROM rag_documents + `) + if err != nil { + return nil, fmt.Errorf("failed to query documents: %w", err) + } + defer rows.Close() + + type scoredDoc struct { + doc rag.Document + score float64 + } + + var docs []scoredDoc + + for rows.Next() { + // Check for context cancellation during iteration + select { + case <-ctx.Done(): + return nil, ctx.Err() + default: + } + + var title, url, content string + var embeddingBytes []byte + + if err := rows.Scan(&title, &url, &content, &embeddingBytes); err != nil { + s.log.WithError(err).Warn("failed to scan row") + continue + } + + // Parse embedding from stored bytes + docEmbedding, err := parseEmbedding(embeddingBytes) + if err != nil { + s.log.WithError(err).Warn("failed to parse embedding") + continue + } + + // Calculate cosine similarity + score := cosineSimilarity(queryEmbedding, docEmbedding) + + docs = append(docs, scoredDoc{ + doc: rag.Document{ + Title: title, + URL: url, + Content: content, + Score: score, + }, + score: score, + }) + } + + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("error iterating rows: %w", err) + } + + // Sort by score descending + sort.Slice(docs, func(i, j int) bool { + return docs[i].score > docs[j].score + }) + + // Take top-k results + result := &rag.SearchResult{ + Documents: make([]rag.Document, 0, maxResults), + } + + for i := 0; i < len(docs) && i < maxResults; i++ { + result.Documents = append(result.Documents, docs[i].doc) + } + + s.log.WithFields(logrus.Fields{ + "query": query, + "total_docs": len(docs), + "returned": len(result.Documents), + "max_results": maxResults, + }).Debug("search completed") + + return result, nil +} + +// Close cleans up resources. +func (s *Searcher) Close() error { + return s.db.Close() +} + +// embed returns the embedding vector for the given text using Vertex AI. +func (s *Searcher) embed(ctx context.Context, text string) ([]float32, error) { + contents := []*genai.Content{ + {Parts: []*genai.Part{{Text: text}}}, + } + + result, err := s.genaiClient.Models.EmbedContent(ctx, s.embeddingModel, contents, nil) + if err != nil { + return nil, fmt.Errorf("failed to embed text: %w", err) + } + + if result == nil || len(result.Embeddings) == 0 { + return nil, fmt.Errorf("no embeddings returned") + } + + return result.Embeddings[0].Values, nil +} + +// parseEmbedding converts stored bytes to a float32 slice. +// The embedding is stored as a binary blob of float32 values. +func parseEmbedding(data []byte) ([]float32, error) { + if len(data) == 0 { + return nil, fmt.Errorf("empty embedding data") + } + + // Each float32 is 4 bytes + if len(data)%4 != 0 { + return nil, fmt.Errorf("invalid embedding data length: %d", len(data)) + } + + numFloats := len(data) / 4 + embedding := make([]float32, numFloats) + + for i := 0; i < numFloats; i++ { + offset := i * 4 + bits := uint32(data[offset]) | + uint32(data[offset+1])<<8 | + uint32(data[offset+2])<<16 | + uint32(data[offset+3])<<24 + embedding[i] = math.Float32frombits(bits) + } + + return embedding, nil +} + +// cosineSimilarity calculates the cosine similarity between two vectors. +func cosineSimilarity(a, b []float32) float64 { + if len(a) != len(b) || len(a) == 0 { + return 0 + } + + var dotProduct, normA, normB float64 + + for i := range a { + dotProduct += float64(a[i]) * float64(b[i]) + normA += float64(a[i]) * float64(a[i]) + normB += float64(b[i]) * float64(b[i]) + } + + if normA == 0 || normB == 0 { + return 0 + } + + return dotProduct / (math.Sqrt(normA) * math.Sqrt(normB)) +} diff --git a/internal/agent/rag/rag.go b/internal/agent/rag/rag.go new file mode 100644 index 000000000..943000626 --- /dev/null +++ b/internal/agent/rag/rag.go @@ -0,0 +1,42 @@ +package rag + +import "context" + +// DocumentSearcher defines the interface for RAG document retrieval. +// Implementations can use DuckDB, Elasticsearch, or other providers. +type DocumentSearcher interface { + // Search queries the document store and returns relevant excerpts. + Search(ctx context.Context, query string, opts *SearchOptions) (*SearchResult, error) + + // Close cleans up any resources held by the searcher. + Close() error +} + +// SearchOptions configures the document search behavior. +type SearchOptions struct { + // MaxResults is the maximum number of documents to return. + MaxResults int + + // Filters allows filtering by metadata (e.g., {"source": "docs.nais.io"}). + Filters map[string]string +} + +// SearchResult contains the documents returned from a search. +type SearchResult struct { + Documents []Document +} + +// Document represents a retrieved document excerpt. +type Document struct { + // Title is the document or section title. + Title string + + // Content is the relevant text excerpt. + Content string + + // URL is the source URL for the document. + URL string + + // Score indicates the relevance score (higher is more relevant). + Score float64 +} diff --git a/internal/agent/tools/chart.go b/internal/agent/tools/chart.go new file mode 100644 index 000000000..ebf61baee --- /dev/null +++ b/internal/agent/tools/chart.go @@ -0,0 +1,56 @@ +// Package tools provides tool definitions and execution for the agent. +package tools + +import ( + "context" + "fmt" +) + +// ChartTools provides tools for chart visualization. +type ChartTools struct{} + +// NewChartTools creates a new ChartTools instance. +func NewChartTools() *ChartTools { + return &ChartTools{} +} + +// RenderChart validates chart parameters and prepares the chart data. +// Note: The actual rendering happens on the client side using the returned ChartData. +func (t *ChartTools) RenderChart(ctx context.Context, input ChartData) (*ChartData, error) { + // Validate required fields + if input.ChartType == "" { + return nil, fmt.Errorf("chart_type is required") + } + if input.ChartType != "line" { + return nil, fmt.Errorf("unsupported chart_type: %s (only 'line' is supported)", input.ChartType) + } + + if input.Title == "" { + return nil, fmt.Errorf("title is required") + } + + if input.Environment == "" { + return nil, fmt.Errorf("environment is required") + } + + if input.Query == "" { + return nil, fmt.Errorf("query is required") + } + + // Validate optional fields + if input.Interval != "" { + validIntervals := map[string]bool{"1h": true, "6h": true, "1d": true, "7d": true, "30d": true} + if !validIntervals[input.Interval] { + return nil, fmt.Errorf("invalid interval: %s (valid values: 1h, 6h, 1d, 7d, 30d)", input.Interval) + } + } + + if input.YFormat != "" { + validFormats := map[string]bool{"number": true, "percentage": true, "bytes": true, "cpu_cores": true, "duration": true} + if !validFormats[input.YFormat] { + return nil, fmt.Errorf("invalid y_format: %s (valid values: number, percentage, bytes, cpu_cores, duration)", input.YFormat) + } + } + + return &input, nil +} diff --git a/internal/agent/tools/chart_test.go b/internal/agent/tools/chart_test.go new file mode 100644 index 000000000..ae720469b --- /dev/null +++ b/internal/agent/tools/chart_test.go @@ -0,0 +1,252 @@ +package tools + +import ( + "context" + "testing" +) + +func TestRenderChart(t *testing.T) { + ct := NewChartTools() + ctx := context.Background() + + tests := []struct { + name string + input ChartData + wantErr bool + errContains string + validate func(*testing.T, *ChartData) + }{ + { + name: "valid chart with all fields", + input: ChartData{ + ChartType: "line", + Title: "CPU Usage", + Environment: "dev", + Query: "sum(rate(container_cpu_usage_seconds_total[5m]))", + Interval: "1h", + YFormat: "cpu_cores", + LabelTemplate: "{pod}/{container}", + }, + wantErr: false, + validate: func(t *testing.T, c *ChartData) { + if c.ChartType != "line" { + t.Errorf("expected chart_type 'line', got %q", c.ChartType) + } + if c.Title != "CPU Usage" { + t.Errorf("expected title 'CPU Usage', got %q", c.Title) + } + if c.Environment != "dev" { + t.Errorf("expected environment 'dev', got %q", c.Environment) + } + if c.Query != "sum(rate(container_cpu_usage_seconds_total[5m]))" { + t.Errorf("unexpected query: %q", c.Query) + } + if c.Interval != "1h" { + t.Errorf("expected interval '1h', got %q", c.Interval) + } + if c.YFormat != "cpu_cores" { + t.Errorf("expected y_format 'cpu_cores', got %q", c.YFormat) + } + if c.LabelTemplate != "{pod}/{container}" { + t.Errorf("expected label_template '{pod}/{container}', got %q", c.LabelTemplate) + } + }, + }, + { + name: "valid chart with required fields only", + input: ChartData{ + ChartType: "line", + Title: "Memory Usage", + Environment: "prod", + Query: "container_memory_usage_bytes", + }, + wantErr: false, + validate: func(t *testing.T, c *ChartData) { + if c.ChartType != "line" { + t.Errorf("expected chart_type 'line', got %q", c.ChartType) + } + if c.Interval != "" { + t.Errorf("expected empty interval, got %q", c.Interval) + } + if c.YFormat != "" { + t.Errorf("expected empty y_format, got %q", c.YFormat) + } + if c.LabelTemplate != "" { + t.Errorf("expected empty label_template, got %q", c.LabelTemplate) + } + }, + }, + { + name: "missing chart_type", + input: ChartData{Title: "CPU Usage", Environment: "dev", Query: "some_query"}, + wantErr: true, + errContains: "chart_type is required", + }, + { + name: "empty chart_type", + input: ChartData{ChartType: "", Title: "CPU Usage", Environment: "dev", Query: "some_query"}, + wantErr: true, + errContains: "chart_type is required", + }, + { + name: "unsupported chart_type", + input: ChartData{ChartType: "bar", Title: "CPU Usage", Environment: "dev", Query: "some_query"}, + wantErr: true, + errContains: "unsupported chart_type", + }, + { + name: "missing title", + input: ChartData{ChartType: "line", Environment: "dev", Query: "some_query"}, + wantErr: true, + errContains: "title is required", + }, + { + name: "missing environment", + input: ChartData{ChartType: "line", Title: "CPU Usage", Query: "some_query"}, + wantErr: true, + errContains: "environment is required", + }, + { + name: "missing query", + input: ChartData{ChartType: "line", Title: "CPU Usage", Environment: "dev"}, + wantErr: true, + errContains: "query is required", + }, + { + name: "invalid interval", + input: ChartData{ChartType: "line", Title: "CPU Usage", Environment: "dev", Query: "q", Interval: "2h"}, + wantErr: true, + errContains: "invalid interval", + }, + { + name: "valid interval 6h", + input: ChartData{ChartType: "line", Title: "CPU Usage", Environment: "dev", Query: "q", Interval: "6h"}, + validate: func(t *testing.T, c *ChartData) { + if c.Interval != "6h" { + t.Errorf("expected interval '6h', got %q", c.Interval) + } + }, + }, + { + name: "valid interval 1d", + input: ChartData{ChartType: "line", Title: "CPU Usage", Environment: "dev", Query: "q", Interval: "1d"}, + validate: func(t *testing.T, c *ChartData) { + if c.Interval != "1d" { + t.Errorf("expected interval '1d', got %q", c.Interval) + } + }, + }, + { + name: "valid interval 7d", + input: ChartData{ChartType: "line", Title: "CPU Usage", Environment: "dev", Query: "q", Interval: "7d"}, + validate: func(t *testing.T, c *ChartData) { + if c.Interval != "7d" { + t.Errorf("expected interval '7d', got %q", c.Interval) + } + }, + }, + { + name: "valid interval 30d", + input: ChartData{ChartType: "line", Title: "CPU Usage", Environment: "dev", Query: "q", Interval: "30d"}, + validate: func(t *testing.T, c *ChartData) { + if c.Interval != "30d" { + t.Errorf("expected interval '30d', got %q", c.Interval) + } + }, + }, + { + name: "invalid y_format", + input: ChartData{ChartType: "line", Title: "CPU Usage", Environment: "dev", Query: "q", YFormat: "invalid"}, + wantErr: true, + errContains: "invalid y_format", + }, + { + name: "valid y_format number", + input: ChartData{ChartType: "line", Title: "CPU Usage", Environment: "dev", Query: "q", YFormat: "number"}, + validate: func(t *testing.T, c *ChartData) { + if c.YFormat != "number" { + t.Errorf("expected y_format 'number', got %q", c.YFormat) + } + }, + }, + { + name: "valid y_format percentage", + input: ChartData{ChartType: "line", Title: "CPU Usage", Environment: "dev", Query: "q", YFormat: "percentage"}, + validate: func(t *testing.T, c *ChartData) { + if c.YFormat != "percentage" { + t.Errorf("expected y_format 'percentage', got %q", c.YFormat) + } + }, + }, + { + name: "valid y_format bytes", + input: ChartData{ChartType: "line", Title: "CPU Usage", Environment: "dev", Query: "q", YFormat: "bytes"}, + validate: func(t *testing.T, c *ChartData) { + if c.YFormat != "bytes" { + t.Errorf("expected y_format 'bytes', got %q", c.YFormat) + } + }, + }, + { + name: "valid y_format cpu_cores", + input: ChartData{ChartType: "line", Title: "CPU Usage", Environment: "dev", Query: "q", YFormat: "cpu_cores"}, + validate: func(t *testing.T, c *ChartData) { + if c.YFormat != "cpu_cores" { + t.Errorf("expected y_format 'cpu_cores', got %q", c.YFormat) + } + }, + }, + { + name: "valid y_format duration", + input: ChartData{ChartType: "line", Title: "CPU Usage", Environment: "dev", Query: "q", YFormat: "duration"}, + validate: func(t *testing.T, c *ChartData) { + if c.YFormat != "duration" { + t.Errorf("expected y_format 'duration', got %q", c.YFormat) + } + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + chart, err := ct.RenderChart(ctx, tt.input) + + if tt.wantErr { + if err == nil { + t.Errorf("expected error containing %q, got nil", tt.errContains) + return + } + if tt.errContains != "" && !containsStr(err.Error(), tt.errContains) { + t.Errorf("expected error containing %q, got %q", tt.errContains, err.Error()) + } + return + } + + if err != nil { + t.Errorf("unexpected error: %v", err) + return + } + + if chart == nil { + t.Error("expected chart data, got nil") + return + } + + if tt.validate != nil { + tt.validate(t, chart) + } + }) + } +} + +func containsStr(s, substr string) bool { + if len(substr) == 0 { + return true + } + for i := 0; i <= len(s)-len(substr); i++ { + if s[i:i+len(substr)] == substr { + return true + } + } + return false +} diff --git a/internal/agent/tools/graphql.go b/internal/agent/tools/graphql.go new file mode 100644 index 000000000..0719bdec3 --- /dev/null +++ b/internal/agent/tools/graphql.go @@ -0,0 +1,182 @@ +// Package tools provides tool definitions and execution for the agent. +package tools + +import ( + "context" + _ "embed" + "encoding/json" + "fmt" + + "github.com/vektah/gqlparser/v2" + "github.com/vektah/gqlparser/v2/ast" +) + +// NaisAPIGuidance provides context about the API structure and common patterns for LLMs. +// +//go:embed nais_api_guidance.md +var NaisAPIGuidance string + +// GraphQLClient is the interface for executing GraphQL queries. +type GraphQLClient interface { + // ExecuteGraphQL runs a GraphQL query with the given variables. + ExecuteGraphQL(ctx context.Context, query string, variables map[string]any) (map[string]any, error) + + // GetCurrentUser returns information about the current authenticated user. + GetCurrentUser(ctx context.Context) (*UserInfo, error) + + // GetUserTeams returns the teams the current user belongs to. + GetUserTeams(ctx context.Context) ([]TeamInfo, error) +} + +// UserInfo contains information about an authenticated user. +type UserInfo struct { + Name string `json:"name"` + IsAdmin bool `json:"is_admin,omitempty"` +} + +// TeamInfo contains information about a team the user belongs to. +type TeamInfo struct { + Slug string `json:"slug"` + Purpose string `json:"purpose,omitempty"` + Role string `json:"role"` +} + +// GraphQLTools provides GraphQL execution functionality. +type GraphQLTools struct { + client GraphQLClient + schema *ast.Schema + consoleBaseURL string + urlPatterns map[string]string +} + +// NewGraphQLTools creates a new GraphQLTools instance. +func NewGraphQLTools(client GraphQLClient, schema *ast.Schema, consoleBaseURL string, urlPatterns map[string]string) *GraphQLTools { + return &GraphQLTools{ + client: client, + schema: schema, + consoleBaseURL: consoleBaseURL, + urlPatterns: urlPatterns, + } +} + +// GetNaisContext returns the current user, their teams, and console URL information. +func (g *GraphQLTools) GetNaisContext(ctx context.Context) (GetNaisContextOutput, error) { + // Get current user + user, err := g.client.GetCurrentUser(ctx) + if err != nil { + return GetNaisContextOutput{}, fmt.Errorf("failed to get current user: %w", err) + } + + // Get user's teams + teams, err := g.client.GetUserTeams(ctx) + if err != nil { + return GetNaisContextOutput{}, fmt.Errorf("failed to get user teams: %w", err) + } + + return GetNaisContextOutput{ + User: *user, + Teams: teams, + ConsoleBaseURL: g.consoleBaseURL, + ConsoleURLPatterns: g.urlPatterns, + }, nil +} + +// ExecuteGraphQL executes a GraphQL query after validation. +func (g *GraphQLTools) ExecuteGraphQL(ctx context.Context, input ExecuteGraphQLInput) (map[string]any, error) { + variablesStr := input.Variables + if variablesStr == "" { + variablesStr = "{}" + } + + // Parse variables + var variables map[string]any + if err := json.Unmarshal([]byte(variablesStr), &variables); err != nil { + return nil, fmt.Errorf("invalid variables JSON: %w", err) + } + + // Validate the query + validationResult := g.validateQuery(input.Query) + if !validationResult.Valid { + return nil, fmt.Errorf("invalid query: %s", validationResult.Error) + } + + // Execute the query + result, err := g.client.ExecuteGraphQL(ctx, input.Query, variables) + if err != nil { + return nil, fmt.Errorf("query execution failed: %w", err) + } + + return result, nil +} + +// ValidateGraphQL validates a GraphQL query without executing it. +func (g *GraphQLTools) ValidateGraphQL(ctx context.Context, input ValidateGraphQLInput) (ValidateGraphQLOutput, error) { + result := g.validateQuery(input.Query) + + if result.Valid { + return ValidateGraphQLOutput{ + Valid: true, + OperationType: result.OperationType, + OperationName: result.OperationName, + Depth: result.Depth, + }, nil + } + + return ValidateGraphQLOutput{ + Valid: false, + Error: result.Error, + }, nil +} + +// validateQuery validates a GraphQL query against the schema. +func (g *GraphQLTools) validateQuery(query string) *QueryValidationResult { + // Parse the query against the schema + doc, errList := gqlparser.LoadQuery(g.schema, query) + if len(errList) > 0 { + return &QueryValidationResult{ + Valid: false, + Error: errList.Error(), + } + } + + // Check that we have at least one operation + if len(doc.Operations) == 0 { + return &QueryValidationResult{ + Valid: false, + Error: "no operations found in query", + } + } + + // Check operation type - only allow queries + op := doc.Operations[0] + if op.Operation != ast.Query { + return &QueryValidationResult{ + Valid: false, + Error: fmt.Sprintf("only query operations are allowed, got: %s", op.Operation), + } + } + + // Check query depth + depth := calculateQueryDepth(op.SelectionSet, 0) + if depth > maxQueryDepth { + return &QueryValidationResult{ + Valid: false, + Error: fmt.Sprintf("query depth %d exceeds maximum allowed depth of %d", depth, maxQueryDepth), + } + } + + // Check for forbidden secret-related types and fields + if found, reason := checkForSecrets(op.SelectionSet, g.schema); found { + return &QueryValidationResult{ + Valid: false, + Error: reason, + } + } + + return &QueryValidationResult{ + Valid: true, + OperationType: string(op.Operation), + OperationName: op.Name, + Depth: depth, + } +} diff --git a/internal/agent/tools/nais_api_guidance.md b/internal/agent/tools/nais_api_guidance.md new file mode 100644 index 000000000..a69e5be1f --- /dev/null +++ b/internal/agent/tools/nais_api_guidance.md @@ -0,0 +1,163 @@ +## Nais API Guidance + +The Nais API is a GraphQL API for managing applications and jobs on the Nais platform. + +### Key Concepts + +- **Team**: The primary organizational unit. All resources belong to a team. +- **Application**: A long-running workload (deployment) managed by Nais (Often called only App/app). +- **Job**: A scheduled or one-off workload (CronJob/Job) managed by Nais. +- **Environment**: A Kubernetes cluster/namespace where workloads run (e.g., "dev", "prod"). +- **Workload**: A union type representing either an Application or Job. + +### Common Query Patterns + +1. **Get current user and their teams**: + ```graphql + query { me { ... on User { teams(first: 50) { nodes { team { slug } } } } } } + ``` + +2. **Get team details**: + ```graphql + query($slug: Slug!) { team(slug: $slug) { slug purpose slackChannel } } + ``` + +3. **List applications for a team**: + ```graphql + query($slug: Slug!) { + team(slug: $slug) { + applications(first: 50) { + nodes { name state teamEnvironment { environment { name } } } + } + } + } + ``` + +4. **Get application details with instances**: + ```graphql + query($slug: Slug!, $name: String!, $env: [String!]) { + team(slug: $slug) { + applications(filter: { name: $name, environments: $env }, first: 1) { + nodes { + name state + instances { nodes { name restarts status { state message } } } + image { name tag } + } + } + } + } + ``` + +5. **List jobs for a team**: + ```graphql + query($slug: Slug!) { + team(slug: $slug) { + jobs(first: 50) { + nodes { name state schedule { expression } teamEnvironment { environment { name } } } + } + } + } + ``` + +6. **Get vulnerabilities for a workload**: + ```graphql + query($slug: Slug!, $name: String!, $env: [String!]) { + team(slug: $slug) { + applications(filter: { name: $name, environments: $env }, first: 1) { + nodes { + image { + vulnerabilitySummary { critical high medium low } + vulnerabilities(first: 20) { nodes { identifier severity package } } + } + } + } + } + } + ``` + +7. **Get cost information**: + ```graphql + query($slug: Slug!) { + team(slug: $slug) { + cost { monthlySummary { sum } } + environments { + environment { name } + cost { daily(from: "2024-01-01", to: "2024-01-31") { sum } } + } + } + } + ``` + +8. **Search across resources**: + ```graphql + query($query: String!) { + search(filter: { query: $query }, first: 20) { + nodes { + __typename + ... on Application { name team { slug } } + ... on Job { name team { slug } } + ... on Team { slug } + } + } + } + ``` + +9. **Get alerts for a team**: + ```graphql + query($slug: Slug!) { + team(slug: $slug) { + alerts(first: 50) { + nodes { name state teamEnvironment { environment { name } } } + } + } + } + ``` + +10. **Get deployments**: + ```graphql + query($slug: Slug!) { + team(slug: $slug) { + deployments(first: 20) { + nodes { createdAt repository commitSha statuses { nodes { state } } } + } + } + } + ``` + +### Important Types + +- **Slug**: A string identifier for teams (e.g., "my-team") +- **Cursor**: Used for pagination (pass to "after" argument) +- **Date**: Format "YYYY-MM-DD" for cost queries +- **ApplicationState**: RUNNING, NOT_RUNNING, UNKNOWN +- **JobState**: RUNNING, NOT_RUNNING, UNKNOWN +- **Severity**: CRITICAL, HIGH, MEDIUM, LOW, UNASSIGNED (for issues/vulnerabilities) + +### Pagination + +Most list fields support pagination with: +- `first: Int` - Number of items to fetch +- `after: Cursor` - Cursor from previous page +- `pageInfo { hasNextPage endCursor totalCount }` - Pagination info + +### Filtering + +Many fields support filters: +- `filter: { name: String, environments: [String!] }` - For applications/jobs +- `filter: { severity: Severity }` - For issues + +### Tips + +1. DO NOT query secret-related types/fields (Secret, SecretValue, etc.) +2. Always use `__typename` when querying union/interface types (Workload, Issue, etc.) +3. Use fragment spreads for type-specific fields: `... on Application { ingresses { url } }` +4. Start with schema exploration to discover available fields +5. Use pagination for large result sets (default to first: 50) + +### Nais Console URLs + +When providing links to the user, use the console URL patterns provided by the `get_nais_context` tool. +Call `get_nais_context` to get the base URL and all available URL patterns with placeholders. +Replace the placeholders (e.g., `{team}`, `{env}`, `{app}`) with actual values from query results. + +**Note**: Do NOT invent or guess URLs. Only use the URL patterns from `get_nais_context` with actual data from query results. \ No newline at end of file diff --git a/internal/agent/tools/registry.go b/internal/agent/tools/registry.go new file mode 100644 index 000000000..d66bb49bd --- /dev/null +++ b/internal/agent/tools/registry.go @@ -0,0 +1,366 @@ +// Package tools provides tool definitions and execution for the agent. +// This is a simplified implementation for direct LLM integration without MCP protocol overhead. +package tools + +import ( + "context" + "encoding/json" + "fmt" + "log/slog" + + "github.com/vektah/gqlparser/v2/ast" +) + +// Tool represents a tool that can be called by the LLM. +type Tool struct { + Name string + Description string + Parameters []Parameter + Handler Handler +} + +// Parameter describes a tool parameter. +type Parameter struct { + Name string + Type string + Description string + Required bool +} + +// Handler is the function signature for tool handlers. +type Handler func(ctx context.Context, args map[string]any) (any, error) + +// Registry holds all available tools. +type Registry struct { + tools map[string]Tool + toolOrder []string // Maintains registration order + schemaTools *SchemaTools + graphqlTools *GraphQLTools + chartTools *ChartTools + logger *slog.Logger +} + +// RegistryConfig holds configuration for creating a Registry. +type RegistryConfig struct { + // Client is the GraphQL client for executing queries. + Client GraphQLClient + + // Schema is the parsed GraphQL schema. + Schema *ast.Schema + + // ConsoleBaseURL is the base URL for the Nais console. + ConsoleBaseURL string + + // ConsoleURLPatterns are URL patterns for console pages. + ConsoleURLPatterns map[string]string + + // Logger is the logger for tool operations. + Logger *slog.Logger +} + +// NewRegistry creates a new tool registry with all tools registered. +func NewRegistry(cfg RegistryConfig) *Registry { + if cfg.Logger == nil { + cfg.Logger = slog.Default() + } + + r := &Registry{ + tools: make(map[string]Tool), + toolOrder: make([]string, 0), + schemaTools: NewSchemaTools(cfg.Schema), + graphqlTools: NewGraphQLTools(cfg.Client, cfg.Schema, cfg.ConsoleBaseURL, cfg.ConsoleURLPatterns), + chartTools: NewChartTools(), + logger: cfg.Logger, + } + + r.registerAllTools() + return r +} + +// register adds a tool to the registry. +func (r *Registry) register(tool Tool) { + r.tools[tool.Name] = tool + r.toolOrder = append(r.toolOrder, tool.Name) +} + +// registerAllTools registers all available tools. +func (r *Registry) registerAllTools() { + // Schema exploration tools + r.register(Tool{ + Name: "schema_list_types", + Description: "List all types in the Nais GraphQL API schema, grouped by kind. Use this to explore available data types before querying specific type details.", + Parameters: []Parameter{ + {Name: "kind", Type: "string", Description: "Filter by kind: 'OBJECT', 'INTERFACE', 'ENUM', 'UNION', 'INPUT_OBJECT', 'SCALAR', or 'all' (default: 'all')", Required: false}, + {Name: "search", Type: "string", Description: "Filter type names containing this string (case-insensitive)", Required: false}, + }, + Handler: r.executeSchemaListTypes, + }) + + r.register(Tool{ + Name: "schema_get_type", + Description: "Get complete details about a GraphQL type: fields with their types, interfaces it implements, types that implement it (for interfaces), enum values, or union member types.", + Parameters: []Parameter{ + {Name: "name", Type: "string", Description: "The exact type name (e.g., 'Application', 'Team', 'DeploymentState')", Required: true}, + }, + Handler: r.executeSchemaGetType, + }) + + r.register(Tool{ + Name: "schema_list_queries", + Description: "List all available GraphQL query operations with their return types and number of arguments. These are the entry points for reading data from the Nais API.", + Parameters: []Parameter{ + {Name: "search", Type: "string", Description: "Filter query names or descriptions containing this string (case-insensitive)", Required: false}, + }, + Handler: r.executeSchemaListQueries, + }) + + r.register(Tool{ + Name: "schema_list_mutations", + Description: "List all available GraphQL mutation operations with their return types and number of arguments. Mutations are used to modify data (note: the agent currently only exposes read operations).", + Parameters: []Parameter{ + {Name: "search", Type: "string", Description: "Filter mutation names or descriptions containing this string (case-insensitive)", Required: false}, + }, + Handler: r.executeSchemaListMutations, + }) + + r.register(Tool{ + Name: "schema_get_field", + Description: "Get detailed information about a specific field including its arguments with types and defaults, return type, description, and deprecation status. Use 'Query' as the type to inspect query operations, or 'Mutation' for mutations.", + Parameters: []Parameter{ + {Name: "type", Type: "string", Description: "The type name containing the field (use 'Query' for root queries, 'Mutation' for root mutations, or any object type name)", Required: true}, + {Name: "field", Type: "string", Description: "The field name to inspect", Required: true}, + }, + Handler: r.executeSchemaGetField, + }) + + r.register(Tool{ + Name: "schema_get_enum", + Description: "Get all possible values for an enum type with their descriptions and deprecation status. Use this to understand valid values for enum fields (e.g., ApplicationState, DeploymentState).", + Parameters: []Parameter{ + {Name: "name", Type: "string", Description: "The enum type name (e.g., 'ApplicationState', 'TeamRole')", Required: true}, + }, + Handler: r.executeSchemaGetEnum, + }) + + r.register(Tool{ + Name: "schema_search", + Description: "Search across all schema types, fields, and enum values by name or description. Returns up to 50 matches. Use this to discover relevant types when you're not sure of exact names.", + Parameters: []Parameter{ + {Name: "query", Type: "string", Description: "Search term to match against names and descriptions (case-insensitive)", Required: true}, + }, + Handler: r.executeSchemaSearch, + }) + + r.register(Tool{ + Name: "schema_get_implementors", + Description: "Get all concrete types that implement a GraphQL interface. Use this to find all possible types when a query returns an interface type.", + Parameters: []Parameter{ + {Name: "interface", Type: "string", Description: "The interface name (e.g., 'Workload', 'Issue')", Required: true}, + }, + Handler: r.executeSchemaGetImplementors, + }) + + r.register(Tool{ + Name: "schema_get_union_types", + Description: "Get all member types of a GraphQL union. Use this to understand what concrete types can be returned when a query returns a union type.", + Parameters: []Parameter{ + {Name: "union", Type: "string", Description: "The union type name", Required: true}, + }, + Handler: r.executeSchemaGetUnionTypes, + }) + + // GraphQL execution tools + r.register(Tool{ + Name: "get_nais_context", + Description: "Get the current Nais context including authenticated user, their teams, and console URL. Call this first to understand what the user has access to and to get the correct console URL for links.", + Parameters: []Parameter{}, + Handler: r.executeGetNaisContext, + }) + + r.register(Tool{ + Name: "execute_graphql", + Description: "Execute a GraphQL query against the Nais API.\n\nIMPORTANT: Before using this tool, use the schema exploration tools (schema_list_queries, schema_get_type, schema_get_field) to understand the available types and fields.\n\nThis tool only supports queries (read operations). Mutations are not allowed.\n\n" + NaisAPIGuidance, + Parameters: []Parameter{ + {Name: "query", Type: "string", Description: "The GraphQL query to execute. Must be a query operation (not mutation or subscription).", Required: true}, + {Name: "variables", Type: "string", Description: "JSON object containing variables for the query. Example: {\"slug\": \"my-team\", \"first\": 10}", Required: false}, + }, + Handler: r.executeGraphQL, + }) + + r.register(Tool{ + Name: "validate_graphql", + Description: "Validate a GraphQL query against the schema without executing it. Use this to check if your query is valid before executing.", + Parameters: []Parameter{ + {Name: "query", Type: "string", Description: "The GraphQL query to validate.", Required: true}, + }, + Handler: r.executeValidateGraphQL, + }) + + // Visualization tools + r.register(Tool{ + Name: "render_chart", + Description: `Render a Prometheus metrics chart in the chat. Use this tool when the user asks about metrics, resource usage, trends, or any data that would be better visualized as a chart rather than described in text. + +Currently only line charts are supported. The chart will be rendered by the client using the provided Prometheus query. + +Guidelines for when to use this tool: +- CPU, memory, or network usage over time +- Request rates, error rates, or latency trends +- Any time-series metrics the user wants to visualize +- When comparing metrics across pods or containers + +Do NOT use this tool for: +- Simple numeric values that don't need visualization +- Non-metrics questions +- When the user explicitly asks for text/numbers only`, + Parameters: []Parameter{ + {Name: "chart_type", Type: "string", Description: "The type of chart to render. Currently only 'line' is supported.", Required: true}, + {Name: "title", Type: "string", Description: "A human-readable title for the chart, e.g., 'CPU Usage for my-app'", Required: true}, + {Name: "environment", Type: "string", Description: "The environment to query metrics from (e.g., 'dev', 'prod'). Use the environment from the current context if available.", Required: true}, + {Name: "query", Type: "string", Description: "The Prometheus query to execute. Must be a valid PromQL query.", Required: true}, + {Name: "interval", Type: "string", Description: "Time interval for the query. Valid values: '1h' (1 hour), '6h' (6 hours), '1d' (1 day), '7d' (7 days), '30d' (30 days). Defaults to '1h'.", Required: false}, + {Name: "y_format", Type: "string", Description: "Format type for Y-axis values. Valid values: 'number', 'percentage', 'bytes', 'cpu_cores', 'duration'. Helps the client format the values appropriately.", Required: false}, + {Name: "label_template", Type: "string", Description: "Template string for formatting series labels. Use {label_name} syntax, e.g., '{pod}' or '{pod}/{container}'. If not provided, default label formatting is used.", Required: false}, + }, + Handler: r.executeRenderChart, + }) +} + +// Execute runs a tool by name with the given arguments. +func (r *Registry) Execute(ctx context.Context, name string, args map[string]any) (any, error) { + tool, ok := r.tools[name] + if !ok { + return nil, fmt.Errorf("unknown tool: %s", name) + } + + r.logger.Debug("executing tool", "name", name) + return tool.Handler(ctx, args) +} + +// ListTools returns all registered tools. +func (r *Registry) ListTools() []Tool { + result := make([]Tool, 0, len(r.toolOrder)) + for _, name := range r.toolOrder { + if tool, ok := r.tools[name]; ok { + result = append(result, tool) + } + } + return result +} + +// Tool handlers - Schema tools + +func (r *Registry) executeSchemaListTypes(ctx context.Context, args map[string]any) (any, error) { + input := SchemaListTypesInput{ + Kind: getString(args, "kind"), + Search: getString(args, "search"), + } + return r.schemaTools.ListTypes(ctx, input) +} + +func (r *Registry) executeSchemaGetType(ctx context.Context, args map[string]any) (any, error) { + input := SchemaGetTypeInput{ + Name: getString(args, "name"), + } + return r.schemaTools.GetType(ctx, input) +} + +func (r *Registry) executeSchemaListQueries(ctx context.Context, args map[string]any) (any, error) { + input := SchemaListQueriesInput{ + Search: getString(args, "search"), + } + return r.schemaTools.ListQueries(ctx, input) +} + +func (r *Registry) executeSchemaListMutations(ctx context.Context, args map[string]any) (any, error) { + input := SchemaListMutationsInput{ + Search: getString(args, "search"), + } + return r.schemaTools.ListMutations(ctx, input) +} + +func (r *Registry) executeSchemaGetField(ctx context.Context, args map[string]any) (any, error) { + input := SchemaGetFieldInput{ + Type: getString(args, "type"), + Field: getString(args, "field"), + } + return r.schemaTools.GetField(ctx, input) +} + +func (r *Registry) executeSchemaGetEnum(ctx context.Context, args map[string]any) (any, error) { + input := SchemaGetEnumInput{ + Name: getString(args, "name"), + } + return r.schemaTools.GetEnum(ctx, input) +} + +func (r *Registry) executeSchemaSearch(ctx context.Context, args map[string]any) (any, error) { + input := SchemaSearchInput{ + Query: getString(args, "query"), + } + return r.schemaTools.Search(ctx, input) +} + +func (r *Registry) executeSchemaGetImplementors(ctx context.Context, args map[string]any) (any, error) { + input := SchemaGetImplementorsInput{ + Interface: getString(args, "interface"), + } + return r.schemaTools.GetImplementors(ctx, input) +} + +func (r *Registry) executeSchemaGetUnionTypes(ctx context.Context, args map[string]any) (any, error) { + input := SchemaGetUnionTypesInput{ + Union: getString(args, "union"), + } + return r.schemaTools.GetUnionTypes(ctx, input) +} + +// Tool handlers - GraphQL tools + +func (r *Registry) executeGetNaisContext(ctx context.Context, args map[string]any) (any, error) { + return r.graphqlTools.GetNaisContext(ctx) +} + +func (r *Registry) executeGraphQL(ctx context.Context, args map[string]any) (any, error) { + input := ExecuteGraphQLInput{ + Query: getString(args, "query"), + Variables: getString(args, "variables"), + } + return r.graphqlTools.ExecuteGraphQL(ctx, input) +} + +func (r *Registry) executeValidateGraphQL(ctx context.Context, args map[string]any) (any, error) { + input := ValidateGraphQLInput{ + Query: getString(args, "query"), + } + return r.graphqlTools.ValidateGraphQL(ctx, input) +} + +// Tool handlers - Visualization tools + +func (r *Registry) executeRenderChart(ctx context.Context, args map[string]any) (any, error) { + input := ChartData{ + ChartType: getString(args, "chart_type"), + Title: getString(args, "title"), + Environment: getString(args, "environment"), + Query: getString(args, "query"), + Interval: getString(args, "interval"), + YFormat: getString(args, "y_format"), + LabelTemplate: getString(args, "label_template"), + } + return r.chartTools.RenderChart(ctx, input) +} + +// getString safely extracts a string from args map. +func getString(args map[string]any, key string) string { + if v, ok := args[key]; ok { + if s, ok := v.(string); ok { + return s + } + // Try JSON number or other types + if b, err := json.Marshal(v); err == nil { + return string(b) + } + } + return "" +} diff --git a/internal/agent/tools/schema.go b/internal/agent/tools/schema.go new file mode 100644 index 000000000..20c0baf11 --- /dev/null +++ b/internal/agent/tools/schema.go @@ -0,0 +1,440 @@ +// Package tools provides tool definitions and execution for the agent. +package tools + +import ( + "context" + "fmt" + "sort" + "strings" + + "github.com/vektah/gqlparser/v2/ast" +) + +// SchemaTools provides schema exploration functionality. +type SchemaTools struct { + schema *ast.Schema +} + +// NewSchemaTools creates a new SchemaTools instance. +func NewSchemaTools(schema *ast.Schema) *SchemaTools { + return &SchemaTools{schema: schema} +} + +// ListTypes lists all types in the schema, grouped by kind. +func (s *SchemaTools) ListTypes(ctx context.Context, input SchemaListTypesInput) (SchemaListTypesOutput, error) { + kind := strings.ToUpper(input.Kind) + if kind == "" { + kind = "ALL" + } + search := strings.ToLower(input.Search) + + var output SchemaListTypesOutput + + for name, def := range s.schema.Types { + // Skip built-in types + if strings.HasPrefix(name, "__") { + continue + } + + // Filter by kind + defKind := string(def.Kind) + if kind != "ALL" && defKind != kind { + continue + } + + // Filter by search + if search != "" && !strings.Contains(strings.ToLower(name), search) { + continue + } + + switch def.Kind { + case ast.Object: + output.Objects = append(output.Objects, name) + case ast.Interface: + output.Interfaces = append(output.Interfaces, name) + case ast.Enum: + output.Enums = append(output.Enums, name) + case ast.Union: + output.Unions = append(output.Unions, name) + case ast.InputObject: + output.InputObjects = append(output.InputObjects, name) + case ast.Scalar: + output.Scalars = append(output.Scalars, name) + } + } + + // Sort all lists + sort.Strings(output.Objects) + sort.Strings(output.Interfaces) + sort.Strings(output.Enums) + sort.Strings(output.Unions) + sort.Strings(output.InputObjects) + sort.Strings(output.Scalars) + + return output, nil +} + +// GetType returns details about a specific type. +func (s *SchemaTools) GetType(ctx context.Context, input SchemaGetTypeInput) (SchemaGetTypeOutput, error) { + def, ok := s.schema.Types[input.Name] + if !ok { + return SchemaGetTypeOutput{}, fmt.Errorf("type %q not found", input.Name) + } + + output := SchemaGetTypeOutput{ + Name: def.Name, + Kind: string(def.Kind), + Description: def.Description, + } + + // Add interfaces this type implements + if len(def.Interfaces) > 0 { + output.Implements = append([]string(nil), def.Interfaces...) + } + + // Add fields for OBJECT, INTERFACE, INPUT_OBJECT + if def.Kind == ast.Object || def.Kind == ast.Interface || def.Kind == ast.InputObject { + output.Fields = formatFields(def.Fields) + } + + // Add enum values for ENUM + if def.Kind == ast.Enum { + for _, v := range def.EnumValues { + value := SchemaEnumValue{ + Name: v.Name, + Description: v.Description, + } + if dep := v.Directives.ForName("deprecated"); dep != nil { + if reason := dep.Arguments.ForName("reason"); reason != nil { + value.Deprecated = reason.Value.Raw + } else { + value.Deprecated = true + } + } + output.Values = append(output.Values, value) + } + } + + // Add types for UNION + if def.Kind == ast.Union { + output.Types = append([]string(nil), def.Types...) + } + + // Add implementedBy for INTERFACE + if def.Kind == ast.Interface { + var implementedBy []string + for typeName, typeDef := range s.schema.Types { + for _, iface := range typeDef.Interfaces { + if iface == input.Name { + implementedBy = append(implementedBy, typeName) + break + } + } + } + sort.Strings(implementedBy) + if len(implementedBy) > 0 { + output.ImplementedBy = implementedBy + } + } + + return output, nil +} + +// ListQueries lists all available query operations. +func (s *SchemaTools) ListQueries(ctx context.Context, input SchemaListQueriesInput) ([]SchemaOperationInfo, error) { + search := strings.ToLower(input.Search) + + queryType := s.schema.Query + if queryType == nil { + return nil, fmt.Errorf("query type not found in schema") + } + + var queries []SchemaOperationInfo + for _, field := range queryType.Fields { + if search == "" || strings.Contains(strings.ToLower(field.Name), search) || strings.Contains(strings.ToLower(field.Description), search) { + queries = append(queries, SchemaOperationInfo{ + Name: field.Name, + ReturnType: field.Type.String(), + Description: truncate(field.Description, 150), + ArgCount: len(field.Arguments), + }) + } + } + + // Sort by name + sort.Slice(queries, func(i, j int) bool { + return queries[i].Name < queries[j].Name + }) + + return queries, nil +} + +// ListMutations lists all available mutation operations. +func (s *SchemaTools) ListMutations(ctx context.Context, input SchemaListMutationsInput) ([]SchemaOperationInfo, error) { + search := strings.ToLower(input.Search) + + mutationType := s.schema.Mutation + if mutationType == nil { + return []SchemaOperationInfo{}, nil + } + + var mutations []SchemaOperationInfo + for _, field := range mutationType.Fields { + if search == "" || strings.Contains(strings.ToLower(field.Name), search) || strings.Contains(strings.ToLower(field.Description), search) { + mutations = append(mutations, SchemaOperationInfo{ + Name: field.Name, + ReturnType: field.Type.String(), + Description: truncate(field.Description, 150), + ArgCount: len(field.Arguments), + }) + } + } + + // Sort by name + sort.Slice(mutations, func(i, j int) bool { + return mutations[i].Name < mutations[j].Name + }) + + return mutations, nil +} + +// GetField returns details about a specific field on a type. +func (s *SchemaTools) GetField(ctx context.Context, input SchemaGetFieldInput) (SchemaGetFieldOutput, error) { + typeDef, ok := s.schema.Types[input.Type] + if !ok { + return SchemaGetFieldOutput{}, fmt.Errorf("type %q not found", input.Type) + } + + var field *ast.FieldDefinition + for _, f := range typeDef.Fields { + if f.Name == input.Field { + field = f + break + } + } + + if field == nil { + return SchemaGetFieldOutput{}, fmt.Errorf("field %q not found on type %q", input.Field, input.Type) + } + + output := SchemaGetFieldOutput{ + Name: field.Name, + Type: field.Type.String(), + Description: field.Description, + } + + // Check for deprecation + if dep := field.Directives.ForName("deprecated"); dep != nil { + if reason := dep.Arguments.ForName("reason"); reason != nil { + output.Deprecated = reason.Value.Raw + } else { + output.Deprecated = true + } + } + + // Add arguments + if len(field.Arguments) > 0 { + for _, arg := range field.Arguments { + argInfo := SchemaArgumentInfo{ + Name: arg.Name, + Type: arg.Type.String(), + Description: arg.Description, + } + if arg.DefaultValue != nil { + argInfo.Default = arg.DefaultValue.String() + } + output.Args = append(output.Args, argInfo) + } + } + + return output, nil +} + +// GetEnum returns details about an enum type. +func (s *SchemaTools) GetEnum(ctx context.Context, input SchemaGetEnumInput) (SchemaGetEnumOutput, error) { + def, ok := s.schema.Types[input.Name] + if !ok || def.Kind != ast.Enum { + return SchemaGetEnumOutput{}, fmt.Errorf("enum %q not found", input.Name) + } + + output := SchemaGetEnumOutput{ + Name: def.Name, + Description: def.Description, + } + + for _, v := range def.EnumValues { + value := SchemaEnumValue{ + Name: v.Name, + Description: v.Description, + } + if dep := v.Directives.ForName("deprecated"); dep != nil { + if reason := dep.Arguments.ForName("reason"); reason != nil { + value.Deprecated = reason.Value.Raw + } else { + value.Deprecated = true + } + } + output.Values = append(output.Values, value) + } + + return output, nil +} + +// Search searches across all schema types, fields, and enum values. +func (s *SchemaTools) Search(ctx context.Context, input SchemaSearchInput) (SchemaSearchOutput, error) { + query := strings.ToLower(input.Query) + var results []SchemaSearchResult + + // Search types + for name, def := range s.schema.Types { + // Skip built-in types + if strings.HasPrefix(name, "__") { + continue + } + + // Match type name or description + if strings.Contains(strings.ToLower(name), query) || strings.Contains(strings.ToLower(def.Description), query) { + results = append(results, SchemaSearchResult{ + Kind: strings.ToLower(string(def.Kind)), + Name: name, + Description: truncate(def.Description, 100), + }) + } + + // Search fields + for _, field := range def.Fields { + if strings.Contains(strings.ToLower(field.Name), query) || strings.Contains(strings.ToLower(field.Description), query) { + results = append(results, SchemaSearchResult{ + Kind: "field", + Type: name, + Name: field.Name, + FieldType: field.Type.String(), + Description: truncate(field.Description, 100), + }) + } + } + + // Search enum values + for _, v := range def.EnumValues { + if strings.Contains(strings.ToLower(v.Name), query) || strings.Contains(strings.ToLower(v.Description), query) { + results = append(results, SchemaSearchResult{ + Kind: "enum_value", + Enum: name, + Name: v.Name, + Description: truncate(v.Description, 100), + }) + } + } + } + + // Sort results by kind then name + sort.Slice(results, func(i, j int) bool { + if results[i].Kind != results[j].Kind { + return results[i].Kind < results[j].Kind + } + return results[i].Name < results[j].Name + }) + + // Limit results + totalMatches := len(results) + if totalMatches > 50 { + results = results[:50] + } + + return SchemaSearchOutput{ + TotalMatches: totalMatches, + Results: results, + }, nil +} + +// GetImplementors returns all types that implement a given interface. +func (s *SchemaTools) GetImplementors(ctx context.Context, input SchemaGetImplementorsInput) (SchemaGetImplementorsOutput, error) { + // Verify the interface exists + def, ok := s.schema.Types[input.Interface] + if !ok || def.Kind != ast.Interface { + return SchemaGetImplementorsOutput{}, fmt.Errorf("interface %q not found", input.Interface) + } + + var implementors []SchemaImplementorInfo + for typeName, typeDef := range s.schema.Types { + for _, iface := range typeDef.Interfaces { + if iface == input.Interface { + implementors = append(implementors, SchemaImplementorInfo{ + Name: typeName, + Description: truncate(typeDef.Description, 100), + }) + break + } + } + } + + // Sort by name + sort.Slice(implementors, func(i, j int) bool { + return implementors[i].Name < implementors[j].Name + }) + + return SchemaGetImplementorsOutput{ + Interface: input.Interface, + Description: def.Description, + Implementors: implementors, + Count: len(implementors), + }, nil +} + +// GetUnionTypes returns all member types of a union. +func (s *SchemaTools) GetUnionTypes(ctx context.Context, input SchemaGetUnionTypesInput) (SchemaGetUnionTypesOutput, error) { + def, ok := s.schema.Types[input.Union] + if !ok || def.Kind != ast.Union { + return SchemaGetUnionTypesOutput{}, fmt.Errorf("union %q not found", input.Union) + } + + var types []SchemaUnionMember + for _, typeName := range def.Types { + typeDef := s.schema.Types[typeName] + member := SchemaUnionMember{ + Name: typeName, + } + if typeDef != nil { + member.Description = truncate(typeDef.Description, 100) + } + types = append(types, member) + } + + return SchemaGetUnionTypesOutput{ + Union: input.Union, + Description: def.Description, + Types: types, + Count: len(types), + }, nil +} + +// formatFields converts AST fields to SchemaFieldInfo slice. +func formatFields(fields ast.FieldList) []SchemaFieldInfo { + var result []SchemaFieldInfo + for _, f := range fields { + field := SchemaFieldInfo{ + Name: f.Name, + Type: f.Type.String(), + Description: truncate(f.Description, 150), + } + if dep := f.Directives.ForName("deprecated"); dep != nil { + if reason := dep.Arguments.ForName("reason"); reason != nil { + field.Deprecated = reason.Value.Raw + } else { + field.Deprecated = true + } + } + result = append(result, field) + } + return result +} + +// truncate truncates a string to the specified length. +func truncate(s string, length int) string { + s = strings.ReplaceAll(s, "\n", " ") + s = strings.TrimSpace(s) + if len(s) <= length { + return s + } + return s[:length] + "..." +} diff --git a/internal/agent/tools/security.go b/internal/agent/tools/security.go new file mode 100644 index 000000000..93a761cc6 --- /dev/null +++ b/internal/agent/tools/security.go @@ -0,0 +1,117 @@ +// Package tools provides tool definitions and execution for the agent. +package tools + +import ( + "fmt" + "strings" + + "github.com/vektah/gqlparser/v2/ast" +) + +const maxQueryDepth = 15 + +// forbiddenTypes are GraphQL types that contain sensitive data and should not be accessible via queries. +// These types and their fields expose secret values that should not be returned to LLMs. +var forbiddenTypes = map[string]bool{ + "Secret": true, // The Secret type contains secret values + "SecretValue": true, // SecretValue contains the actual secret data + "SecretConnection": true, // Connection type that returns Secret nodes + "SecretEdge": true, // Edge type that wraps Secret + "DeploymentKey": true, // Contains the actual deployment key + "CreateServiceAccountTokenPayload": true, // Contains the service account token secret + "ServiceAccountToken": true, // Service account token metadata (but secret field is blocked separately) + "ServiceAccountTokenConnection": true, // Connection type that returns ServiceAccountToken nodes + "ServiceAccountTokenEdge": true, // Edge type that wraps ServiceAccountToken +} + +// checkForSecrets recursively checks if a selection set accesses any forbidden secret-related types. +// It validates against the GraphQL schema to ensure queries don't access Secret or SecretValue types. +func checkForSecrets(selectionSet ast.SelectionSet, schema *ast.Schema) (bool, string) { + for _, selection := range selectionSet { + switch sel := selection.(type) { + case *ast.Field: + // Check if this field's definition exists in the schema + if sel.Definition != nil { + // Check if the field returns a forbidden type + typeName := getBaseTypeName(sel.Definition.Type) + if forbiddenTypes[typeName] { + return true, fmt.Sprintf("security policy: field '%s' returns type '%s' which contains sensitive data that cannot be accessed via this interface. Use the Nais Console or CLI to manage secrets directly.", sel.Name, typeName) + } + } + + // Recursively check nested selections + if len(sel.SelectionSet) > 0 { + if found, reason := checkForSecrets(sel.SelectionSet, schema); found { + return true, reason + } + } + case *ast.InlineFragment: + // Check if the inline fragment is on a forbidden type + if sel.TypeCondition != "" && forbiddenTypes[sel.TypeCondition] { + return true, fmt.Sprintf("security policy: inline fragment on type '%s' which contains sensitive data that cannot be accessed via this interface", sel.TypeCondition) + } + // Check inline fragments recursively + if found, reason := checkForSecrets(sel.SelectionSet, schema); found { + return true, reason + } + case *ast.FragmentSpread: + // Check recursively if the fragment accesses secrets + if sel.Definition != nil { + // Check if the fragment type condition is forbidden + if forbiddenTypes[sel.Definition.TypeCondition] { + return true, fmt.Sprintf("security policy: fragment '%s' is on type '%s' which contains sensitive data that cannot be accessed via this interface", sel.Name, sel.Definition.TypeCondition) + } + + if found, reason := checkForSecrets(sel.Definition.SelectionSet, schema); found { + return true, reason + } + } else if strings.Contains(strings.ToLower(sel.Name), "secret") { + // Fallback heuristic if definition is missing + return true, fmt.Sprintf("security policy: fragment '%s' may access sensitive data that cannot be accessed via this interface", sel.Name) + } + } + } + return false, "" +} + +// getBaseTypeName extracts the base type name from a GraphQL type, removing list and non-null wrappers. +func getBaseTypeName(t *ast.Type) string { + if t.Elem != nil { + return getBaseTypeName(t.Elem) + } + return t.Name() +} + +// calculateQueryDepth calculates the maximum depth of a selection set. +func calculateQueryDepth(selectionSet ast.SelectionSet, currentDepth int) int { + if len(selectionSet) == 0 { + return currentDepth + } + + maxDepth := currentDepth + for _, selection := range selectionSet { + var childDepth int + switch sel := selection.(type) { + case *ast.Field: + childDepth = calculateQueryDepth(sel.SelectionSet, currentDepth+1) + case *ast.InlineFragment: + childDepth = calculateQueryDepth(sel.SelectionSet, currentDepth) + case *ast.FragmentSpread: + if sel.Definition != nil { + childDepth = calculateQueryDepth(sel.Definition.SelectionSet, currentDepth) + } else { + // Fallback if definition is missing + childDepth = currentDepth + 1 + } + } + if childDepth > maxDepth { + maxDepth = childDepth + } + } + return maxDepth +} + +// IsForbiddenType checks if a type name is in the forbidden types list. +func IsForbiddenType(typeName string) bool { + return forbiddenTypes[typeName] +} diff --git a/internal/agent/tools/types.go b/internal/agent/tools/types.go new file mode 100644 index 000000000..1fe720024 --- /dev/null +++ b/internal/agent/tools/types.go @@ -0,0 +1,222 @@ +// Package tools provides tool definitions and execution for the agent. +package tools + +// ============================================================================= +// Schema Tool Types +// ============================================================================= + +// SchemaListTypesInput is the input for the schema_list_types tool. +type SchemaListTypesInput struct { + Kind string `json:"kind,omitempty"` + Search string `json:"search,omitempty"` +} + +// SchemaListTypesOutput is the output for the schema_list_types tool. +type SchemaListTypesOutput struct { + Objects []string `json:"objects,omitempty"` + Interfaces []string `json:"interfaces,omitempty"` + Enums []string `json:"enums,omitempty"` + Unions []string `json:"unions,omitempty"` + InputObjects []string `json:"input_objects,omitempty"` + Scalars []string `json:"scalars,omitempty"` +} + +// SchemaGetTypeInput is the input for the schema_get_type tool. +type SchemaGetTypeInput struct { + Name string `json:"name"` +} + +// SchemaFieldInfo describes a field on a GraphQL type. +type SchemaFieldInfo struct { + Name string `json:"name"` + Type string `json:"type"` + Description string `json:"description,omitempty"` + Deprecated any `json:"deprecated,omitempty"` +} + +// SchemaEnumValue describes a value in a GraphQL enum. +type SchemaEnumValue struct { + Name string `json:"name"` + Description string `json:"description,omitempty"` + Deprecated any `json:"deprecated,omitempty"` +} + +// SchemaGetTypeOutput is the output for the schema_get_type tool. +type SchemaGetTypeOutput struct { + Name string `json:"name"` + Kind string `json:"kind"` + Description string `json:"description,omitempty"` + Implements []string `json:"implements,omitempty"` + Fields []SchemaFieldInfo `json:"fields,omitempty"` + Values []SchemaEnumValue `json:"values,omitempty"` + Types []string `json:"types,omitempty"` + ImplementedBy []string `json:"implementedBy,omitempty"` +} + +// SchemaListQueriesInput is the input for the schema_list_queries tool. +type SchemaListQueriesInput struct { + Search string `json:"search,omitempty"` +} + +// SchemaOperationInfo describes a GraphQL operation (query or mutation). +type SchemaOperationInfo struct { + Name string `json:"name"` + ReturnType string `json:"returnType"` + Description string `json:"description,omitempty"` + ArgCount int `json:"argCount"` +} + +// SchemaListMutationsInput is the input for the schema_list_mutations tool. +type SchemaListMutationsInput struct { + Search string `json:"search,omitempty"` +} + +// SchemaGetFieldInput is the input for the schema_get_field tool. +type SchemaGetFieldInput struct { + Type string `json:"type"` + Field string `json:"field"` +} + +// SchemaArgumentInfo describes an argument on a GraphQL field. +type SchemaArgumentInfo struct { + Name string `json:"name"` + Type string `json:"type"` + Description string `json:"description,omitempty"` + Default string `json:"default,omitempty"` +} + +// SchemaGetFieldOutput is the output for the schema_get_field tool. +type SchemaGetFieldOutput struct { + Name string `json:"name"` + Type string `json:"type"` + Description string `json:"description,omitempty"` + Deprecated any `json:"deprecated,omitempty"` + Args []SchemaArgumentInfo `json:"args,omitempty"` +} + +// SchemaGetEnumInput is the input for the schema_get_enum tool. +type SchemaGetEnumInput struct { + Name string `json:"name"` +} + +// SchemaGetEnumOutput is the output for the schema_get_enum tool. +type SchemaGetEnumOutput struct { + Name string `json:"name"` + Description string `json:"description,omitempty"` + Values []SchemaEnumValue `json:"values"` +} + +// SchemaSearchInput is the input for the schema_search tool. +type SchemaSearchInput struct { + Query string `json:"query"` +} + +// SchemaSearchResult is a single result from a schema search. +type SchemaSearchResult struct { + Kind string `json:"kind"` + Name string `json:"name"` + Type string `json:"type,omitempty"` + Enum string `json:"enum,omitempty"` + FieldType string `json:"fieldType,omitempty"` + Description string `json:"description,omitempty"` +} + +// SchemaSearchOutput is the output for the schema_search tool. +type SchemaSearchOutput struct { + TotalMatches int `json:"totalMatches"` + Results []SchemaSearchResult `json:"results"` +} + +// SchemaGetImplementorsInput is the input for the schema_get_implementors tool. +type SchemaGetImplementorsInput struct { + Interface string `json:"interface"` +} + +// SchemaImplementorInfo describes a type that implements an interface. +type SchemaImplementorInfo struct { + Name string `json:"name"` + Description string `json:"description,omitempty"` +} + +// SchemaGetImplementorsOutput is the output for the schema_get_implementors tool. +type SchemaGetImplementorsOutput struct { + Interface string `json:"interface"` + Description string `json:"description,omitempty"` + Implementors []SchemaImplementorInfo `json:"implementors"` + Count int `json:"count"` +} + +// SchemaGetUnionTypesInput is the input for the schema_get_union_types tool. +type SchemaGetUnionTypesInput struct { + Union string `json:"union"` +} + +// SchemaUnionMember describes a member type of a GraphQL union. +type SchemaUnionMember struct { + Name string `json:"name"` + Description string `json:"description,omitempty"` +} + +// SchemaGetUnionTypesOutput is the output for the schema_get_union_types tool. +type SchemaGetUnionTypesOutput struct { + Union string `json:"union"` + Description string `json:"description,omitempty"` + Types []SchemaUnionMember `json:"types"` + Count int `json:"count"` +} + +// ============================================================================= +// GraphQL Tool Types +// ============================================================================= + +// GetNaisContextOutput is the output for the get_nais_context tool. +type GetNaisContextOutput struct { + User UserInfo `json:"user"` + Teams []TeamInfo `json:"teams"` + ConsoleBaseURL string `json:"console_base_url"` + ConsoleURLPatterns map[string]string `json:"console_url_patterns"` +} + +// ExecuteGraphQLInput is the input for the execute_graphql tool. +type ExecuteGraphQLInput struct { + Query string `json:"query"` + Variables string `json:"variables,omitempty"` +} + +// ValidateGraphQLInput is the input for the validate_graphql tool. +type ValidateGraphQLInput struct { + Query string `json:"query"` +} + +// ValidateGraphQLOutput is the output for the validate_graphql tool. +type ValidateGraphQLOutput struct { + Valid bool `json:"valid"` + Error string `json:"error,omitempty"` + OperationType string `json:"operationType,omitempty"` + OperationName string `json:"operationName,omitempty"` + Depth int `json:"depth,omitempty"` +} + +// QueryValidationResult is the internal result of query validation. +type QueryValidationResult struct { + Valid bool + Error string + OperationType string + OperationName string + Depth int +} + +// ============================================================================= +// Visualization Tool Types +// ============================================================================= + +// ChartData represents the data for a rendered chart. +type ChartData struct { + ChartType string `json:"chart_type"` + Title string `json:"title"` + Environment string `json:"environment"` + Query string `json:"query"` + Interval string `json:"interval,omitempty"` + YFormat string `json:"y_format,omitempty"` + LabelTemplate string `json:"label_template,omitempty"` +} diff --git a/internal/cmd/api/api.go b/internal/cmd/api/api.go index ddde77ca7..236c5a734 100644 --- a/internal/cmd/api/api.go +++ b/internal/cmd/api/api.go @@ -14,6 +14,10 @@ import ( aiven_service "github.com/aiven/go-client-codegen" "github.com/joho/godotenv" "github.com/nais/api/internal/activitylog" + "github.com/nais/api/internal/agent" + "github.com/nais/api/internal/agent/chat" + "github.com/nais/api/internal/agent/rag" + "github.com/nais/api/internal/agent/rag/duckdb" "github.com/nais/api/internal/auth/authn" "github.com/nais/api/internal/auth/middleware" "github.com/nais/api/internal/database" @@ -285,6 +289,30 @@ func run(ctx context.Context, cfg *Config, log logrus.FieldLogger) error { return fmt.Errorf("create loki client: %w", err) } + // Initialize agent handler if enabled + var agentHandler *agent.Handler + if cfg.Agent.Enabled { + chatClient, ragClient, err := initAgentClients(ctx, cfg.Agent, log) + if err != nil { + return fmt.Errorf("initializing agent clients: %w", err) + } + // Defer closing clients + defer chatClient.Close() + defer ragClient.Close() + + agentHandler, err = agent.NewHandler(agent.Config{ + ChatClient: chatClient, + RAGClient: ragClient, + GraphQLHandler: graphHandler, + TenantName: cfg.Tenant, + Log: log.WithField("subsystem", "agent"), + }) + if err != nil { + return fmt.Errorf("initializing agent handler: %w", err) + } + log.Info("Agent service enabled") + } + // HTTP server wg.Go(func() error { return runHTTPServer( @@ -313,6 +341,7 @@ func run(ctx context.Context, cfg *Config, log logrus.FieldLogger) error { lokiClient, cfg.AuditLog.ProjectID, cfg.AuditLog.Location, + agentHandler, log.WithField("subsystem", "http"), ) }) @@ -331,7 +360,12 @@ func run(ctx context.Context, cfg *Config, log logrus.FieldLogger) error { }) wg.Go(func() error { - if err := grpc.Run(ctx, cfg.GRPCListenAddress, pool, log.WithField("subsystem", "grpc")); err != nil { + grpcCfg := &grpc.Config{ + ListenAddress: cfg.GRPCListenAddress, + Pool: pool, + Log: log.WithField("subsystem", "grpc"), + } + if err := grpc.Run(ctx, grpcCfg); err != nil { log.WithError(err).Errorf("error in GRPC server") return err } @@ -389,6 +423,7 @@ func run(ctx context.Context, cfg *Config, log logrus.FieldLogger) error { } wg.Go(func() error { + return nil err = issueChecker.RunChecks(ctx) if err != nil { log.WithError(err).Error("running issue checks") @@ -437,3 +472,46 @@ func setupAuthHandler(ctx context.Context, cfg oAuthConfig, log logrus.FieldLogg } return authn.New(cf, log), nil } + +func initAgentClients(ctx context.Context, cfg agentConfig, log logrus.FieldLogger) (chat.StreamingClient, rag.DocumentSearcher, error) { + // Create Vertex AI chat client + chatClient, err := chat.NewClient(ctx, chat.Config{ + ProjectID: cfg.VertexAI.ProjectID, + Location: cfg.VertexAI.Location, + ModelName: cfg.VertexAI.ModelName, + IncludeThoughts: cfg.VertexAI.IncludeThoughts, + }, log.WithField("client", "vertexai-chat")) + if err != nil { + return nil, nil, fmt.Errorf("creating vertexai chat client: %w", err) + } + + // Create RAG client if enabled + var ragClient rag.DocumentSearcher + if cfg.RAGEnabled { + ragClient, err = duckdb.NewSearcher(ctx, duckdb.Config{ + DBPath: cfg.RAG.DuckDBPath, + ProjectID: cfg.VertexAI.ProjectID, + Location: cfg.VertexAI.Location, + EmbeddingModel: cfg.VertexAI.EmbeddingModel, + }, log.WithField("client", "duckdb-rag")) + if err != nil { + chatClient.Close() + return nil, nil, fmt.Errorf("creating duckdb rag client: %w", err) + } + } else { + ragClient = &noopSearcher{} + } + + return chatClient, ragClient, nil +} + +// noopSearcher implements rag.DocumentSearcher but returns no results. +type noopSearcher struct{} + +func (n *noopSearcher) Search(ctx context.Context, query string, opts *rag.SearchOptions) (*rag.SearchResult, error) { + return &rag.SearchResult{Documents: []rag.Document{}}, nil +} + +func (n *noopSearcher) Close() error { + return nil +} diff --git a/internal/cmd/api/config.go b/internal/cmd/api/config.go index f1e8dfad9..9bde211f1 100644 --- a/internal/cmd/api/config.go +++ b/internal/cmd/api/config.go @@ -144,6 +144,7 @@ type Aiven struct { type Config struct { Aiven Aiven + Agent agentConfig // Tenant is the active tenant Tenant string `env:"TENANT,default=dev-nais"` @@ -240,6 +241,48 @@ type pubSubConfig struct { EventsSubscription string `env:"PUBSUB_EVENTS_SUBSCRIPTION,default=nais-api-log-topic-subscription"` } +// agentConfig contains configuration for the AI agent service. +type agentConfig struct { + // Enabled enables the AI agent gRPC service. + Enabled bool `env:"AGENT_ENABLED,default=false"` + + // RAGEnabled enables RAG (retrieval-augmented generation) using DuckDB. + // When disabled, the agent will not search documentation. + RAGEnabled bool `env:"AGENT_RAG_ENABLED,default=true"` + + // VertexAI contains Vertex AI specific configuration for LLM and embeddings. + VertexAI agentVertexAIConfig + + // RAG contains RAG-specific configuration (only used when RAGEnabled is true). + RAG agentRAGConfig +} + +// agentVertexAIConfig contains Vertex AI specific configuration for the agent. +type agentVertexAIConfig struct { + // ProjectID is the GCP project hosting Vertex AI resources. + ProjectID string `env:"AGENT_VERTEX_AI_PROJECT_ID"` + + // Location is the region for Vertex AI (must be in EU, e.g., "europe-west1"). + Location string `env:"AGENT_VERTEX_AI_LOCATION,default=europe-west1"` + + // ModelName is the Gemini model to use for chat. + ModelName string `env:"AGENT_VERTEX_AI_MODEL,default=gemini-2.5-flash"` + + // EmbeddingModel is the model to use for query embeddings. + EmbeddingModel string `env:"AGENT_VERTEX_AI_EMBEDDING_MODEL,default=gemini-embedding-001"` + + // IncludeThoughts enables returning the model's thinking/reasoning in responses. + // When enabled, Gemini 3+ models will include their thought process in the response. + // This is streamed to clients as "thinking" events. + IncludeThoughts bool `env:"AGENT_VERTEX_AI_INCLUDE_THOUGHTS,default=false"` +} + +// agentRAGConfig contains RAG-specific configuration. +type agentRAGConfig struct { + // DuckDBPath is the local path to the DuckDB RAG index file. + DuckDBPath string `env:"AGENT_RAG_DUCKDB_PATH,default=/data/rag_index.duckdb"` +} + // NewConfig creates a new configuration instance from environment variables func NewConfig(ctx context.Context, lookuper envconfig.Lookuper) (*Config, error) { cfg := &Config{} diff --git a/internal/cmd/api/http.go b/internal/cmd/api/http.go index 51f353639..94378c0de 100644 --- a/internal/cmd/api/http.go +++ b/internal/cmd/api/http.go @@ -12,6 +12,7 @@ import ( "github.com/go-chi/chi/v5" "github.com/jackc/pgx/v5/pgxpool" "github.com/nais/api/internal/activitylog" + "github.com/nais/api/internal/agent" "github.com/nais/api/internal/alerts" "github.com/nais/api/internal/auth/authn" "github.com/nais/api/internal/auth/authz" @@ -98,6 +99,7 @@ func runHTTPServer( lokiClient loki.Client, auditLogProjectID string, auditLogLocation string, + agentHandler *agent.Handler, log logrus.FieldLogger, ) error { router := chi.NewRouter() @@ -178,6 +180,33 @@ func runHTTPServer( r.Get("/callback", authHandler.Callback) }) + // Agent chat routes (if enabled) + if agentHandler != nil { + router.Route("/agent", func(r chi.Router) { + middlewares := []func(http.Handler) http.Handler{ + contextDependencies, + } + + if fakes.WithInsecureUserHeader { + middlewares = append(middlewares, middleware.InsecureUserHeader()) + } + + if jwtMiddleware != nil { + middlewares = append(middlewares, jwtMiddleware) + } + + middlewares = append( + middlewares, + middleware.ApiKeyAuthentication(), + middleware.Oauth2Authentication(authHandler), + middleware.RequireAuthenticatedUser(), + ) + r.Use(middlewares...) + agentHandler.RegisterRoutes(r) + }) + log.Info("Agent HTTP routes registered at /agent") + } + srv := &http.Server{ Addr: listenAddress, Handler: router, @@ -333,6 +362,7 @@ func ConfigureGraph( ctx = postgres.NewLoaderContext(ctx, watchers.ZalandoPostgresWatcher, auditLogProjectID, auditLogLocation) ctx = aivencredentials.NewClientContext(ctx, dynamicClients, log) ctx = database.NewLoaderContext(ctx, pool) + ctx = agent.NewLoaderContext(ctx, pool) ctx = issue.NewContext(ctx, pool) ctx = team.NewLoaderContext(ctx, pool, watchers.NamespaceWatcher) ctx = loki.NewLoaderContext(ctx, lokiClient) diff --git a/internal/cmd/gen_rag_index/config.go b/internal/cmd/gen_rag_index/config.go new file mode 100644 index 000000000..deda88953 --- /dev/null +++ b/internal/cmd/gen_rag_index/config.go @@ -0,0 +1,60 @@ +package gen_rag_index + +import ( + "context" + "fmt" + + "github.com/sethvargo/go-envconfig" +) + +// Config holds configuration for the RAG index generator. +type Config struct { + // TenantName is the tenant name used to construct the docs URL. + // The search index is fetched from: https://docs..cloud.nais.io/search/search_index.json + TenantName string `env:"TENANT,required"` + + // OutputPath is the path where the DuckDB file will be written. + OutputPath string `env:"OUTPUT_PATH,default=./data/rag_index.duckdb"` + + // ChunkMaxChars is the maximum number of characters per chunk. + ChunkMaxChars int `env:"CHUNK_MAX_CHARS,default=1500"` + + // VertexAI contains Vertex AI specific configuration for embeddings. + VertexAI VertexAIConfig +} + +// VertexAIConfig contains Vertex AI specific configuration. +type VertexAIConfig struct { + // ProjectID is the GCP project hosting Vertex AI resources. + ProjectID string `env:"AGENT_VERTEX_AI_PROJECT_ID,required"` + + // Location is the region for Vertex AI (must be in EU, e.g., "europe-west1"). + Location string `env:"AGENT_VERTEX_AI_LOCATION,default=europe-west1"` + + // EmbeddingModel is the model to use for generating embeddings. + EmbeddingModel string `env:"AGENT_VERTEX_AI_EMBEDDING_MODEL,default=gemini-embedding-001"` +} + +// NewConfig creates a new configuration instance from environment variables. +func NewConfig(ctx context.Context, lookuper envconfig.Lookuper) (*Config, error) { + cfg := &Config{} + err := envconfig.ProcessWith(ctx, &envconfig.Config{ + Target: cfg, + Lookuper: lookuper, + }) + if err != nil { + return nil, fmt.Errorf("processing config: %w", err) + } + + return cfg, nil +} + +// DocsBaseURL returns the base URL for the documentation site. +func (c *Config) DocsBaseURL() string { + return fmt.Sprintf("https://docs.%s.cloud.nais.io", c.TenantName) +} + +// SearchIndexURL returns the URL for the search index JSON file. +func (c *Config) SearchIndexURL() string { + return fmt.Sprintf("%s/search/search_index.json", c.DocsBaseURL()) +} diff --git a/internal/cmd/gen_rag_index/gen_rag_index.go b/internal/cmd/gen_rag_index/gen_rag_index.go new file mode 100644 index 000000000..429704e08 --- /dev/null +++ b/internal/cmd/gen_rag_index/gen_rag_index.go @@ -0,0 +1,190 @@ +package gen_rag_index + +import ( + "context" + "errors" + "fmt" + "os" + "time" + + "github.com/joho/godotenv" + "github.com/nais/api/internal/agent/chat" + "github.com/nais/api/internal/logger" + "github.com/sethvargo/go-envconfig" + "github.com/sirupsen/logrus" +) + +const ( + exitCodeSuccess = iota + exitCodeConfigError + exitCodeLoggerError + exitCodeRunError +) + +// Run is the main entry point for the RAG index generator. +func Run(ctx context.Context) { + log, err := logger.New("text", "INFO") + if err != nil { + fmt.Printf("logger error: %s\n", err) + os.Exit(exitCodeLoggerError) + } + loadEnvFile(log) + + cfg, err := NewConfig(ctx, envconfig.OsLookuper()) + if err != nil { + log.WithError(err).Error("configuration error") + os.Exit(exitCodeConfigError) + } + + if err := run(ctx, cfg, log); err != nil { + log.WithError(err).Error("fatal error") + os.Exit(exitCodeRunError) + } + + os.Exit(exitCodeSuccess) +} + +func run(ctx context.Context, cfg *Config, log logrus.FieldLogger) error { + startTime := time.Now() + + log.WithFields(logrus.Fields{ + "tenant": cfg.TenantName, + "output": cfg.OutputPath, + "chunk_size": cfg.ChunkMaxChars, + "search_index": cfg.SearchIndexURL(), + }).Info("starting RAG index generation") + + // Step 1: Download search index + log.Info("downloading search index...") + index, err := DownloadSearchIndex(ctx, cfg.SearchIndexURL()) + if err != nil { + return fmt.Errorf("downloading search index: %w", err) + } + log.WithField("docs", len(index.Docs)).Info("downloaded search index") + + // Step 2: Process into pages + log.Info("processing documents into pages...") + pages := ProcessSearchIndex(index, cfg.DocsBaseURL()) + log.WithField("pages", len(pages)).Info("processed pages") + + // Step 3: Chunk pages + log.Info("chunking pages...") + chunks := ChunkPages(pages, cfg.ChunkMaxChars) + log.WithField("chunks", len(chunks)).Info("created chunks") + + // Step 4: Initialize embedding client + log.Info("initializing embedding client...") + embeddingClient, err := chat.NewEmbeddingClient(ctx, chat.EmbeddingConfig{ + ProjectID: cfg.VertexAI.ProjectID, + Location: cfg.VertexAI.Location, + ModelName: cfg.VertexAI.EmbeddingModel, + }, log.WithField("component", "embedding-client")) + if err != nil { + return fmt.Errorf("creating embedding client: %w", err) + } + defer embeddingClient.Close() + + // Step 5: Generate embeddings + log.Info("generating embeddings...") + chunksWithEmbeddings, err := generateEmbeddings(ctx, embeddingClient, chunks, log) + if err != nil { + return fmt.Errorf("generating embeddings: %w", err) + } + log.WithField("embedded", len(chunksWithEmbeddings)).Info("generated embeddings") + + // Step 6: Write to DuckDB + log.Info("writing to DuckDB...") + writer, err := NewWriter(cfg.OutputPath, log.WithField("component", "writer")) + if err != nil { + return fmt.Errorf("creating writer: %w", err) + } + defer writer.Close() + + if err := writer.WriteChunks(ctx, chunksWithEmbeddings); err != nil { + return fmt.Errorf("writing chunks: %w", err) + } + + if err := writer.CreateIndex(ctx); err != nil { + return fmt.Errorf("creating index: %w", err) + } + + // Step 7: Print stats + stats, err := writer.GetStats(ctx) + if err != nil { + return fmt.Errorf("getting stats: %w", err) + } + + duration := time.Since(startTime) + log.WithFields(logrus.Fields{ + "total_chunks": stats.TotalChunks, + "unique_urls": stats.UniqueURLs, + "avg_content_length": fmt.Sprintf("%.0f", stats.AvgContentLength), + "duration": duration.Round(time.Second).String(), + "output": cfg.OutputPath, + }).Info("RAG index generation complete") + + return nil +} + +// generateEmbeddings generates embeddings for all chunks using batch API calls. +// The Vertex AI API supports up to 250 texts per batch, but we use 100 for safety. +func generateEmbeddings(ctx context.Context, client *chat.EmbeddingClient, chunks []Chunk, log logrus.FieldLogger) ([]ChunkWithEmbedding, error) { + const batchSize = 100 + + result := make([]ChunkWithEmbedding, 0, len(chunks)) + totalBatches := (len(chunks) + batchSize - 1) / batchSize + + for batchNum := 0; batchNum < totalBatches; batchNum++ { + start := batchNum * batchSize + end := start + batchSize + if end > len(chunks) { + end = len(chunks) + } + + batch := chunks[start:end] + + // Extract content for embedding + texts := make([]string, len(batch)) + for i, chunk := range batch { + texts[i] = chunk.Content + } + + // Generate embeddings for the batch + embeddings, err := client.EmbedBatch(ctx, texts) + if err != nil { + return nil, fmt.Errorf("embedding batch %d (chunks %d-%d): %w", batchNum+1, start, end-1, err) + } + + // Combine chunks with their embeddings + for i, chunk := range batch { + result = append(result, ChunkWithEmbedding{ + Chunk: chunk, + Embedding: embeddings[i], + }) + } + + // Progress logging + log.WithFields(logrus.Fields{ + "batch": fmt.Sprintf("%d/%d", batchNum+1, totalBatches), + "chunks": fmt.Sprintf("%d/%d", end, len(chunks)), + "percent": fmt.Sprintf("%.1f%%", float64(end)/float64(len(chunks))*100), + }).Info("embedding progress") + } + + return result, nil +} + +// loadEnvFile will load a .env file if it exists. This is useful for local development. +func loadEnvFile(log logrus.FieldLogger) error { + if _, err := os.Stat(".env"); errors.Is(err, os.ErrNotExist) { + log.Infof("no .env file found") + return nil + } + + if err := godotenv.Load(".env"); err != nil { + return err + } + + log.Infof("loaded .env file") + return nil +} diff --git a/internal/cmd/gen_rag_index/process.go b/internal/cmd/gen_rag_index/process.go new file mode 100644 index 000000000..4cb163aa0 --- /dev/null +++ b/internal/cmd/gen_rag_index/process.go @@ -0,0 +1,387 @@ +package gen_rag_index + +import ( + "context" + "encoding/json" + "fmt" + "html" + "io" + "net/http" + "regexp" + "strings" +) + +// SearchIndex represents the MkDocs search index structure. +type SearchIndex struct { + Config SearchConfig `json:"config"` + Docs []SearchDoc `json:"docs"` +} + +// SearchConfig contains search configuration from MkDocs. +type SearchConfig struct { + Lang []string `json:"lang"` + Separator string `json:"separator"` + Pipeline []string `json:"pipeline"` + Fields map[string]FieldConfig `json:"fields"` +} + +// FieldConfig contains field-specific search configuration. +type FieldConfig struct { + Boost float64 `json:"boost"` +} + +// SearchDoc represents a single document in the search index. +type SearchDoc struct { + Location string `json:"location"` + Title string `json:"title"` + Text string `json:"text"` + Tags []string `json:"tags,omitempty"` +} + +// Page represents a processed documentation page. +type Page struct { + // Path is the page path (location without fragment). + Path string + + // Title is the page title. + Title string + + // Tags are the tags associated with this page. + Tags []string + + // Sections contains the page sections. + Sections []Section + + // URL is the full URL to the page. + URL string +} + +// Section represents a section within a page. +type Section struct { + Title string + Text string +} + +// Chunk represents a chunk of content ready for embedding. +type Chunk struct { + // Title is the page title (included in all chunks from the same page). + Title string + + // URL is the full URL to the page. + URL string + + // Content is the formatted content for embedding. + Content string +} + +// DownloadSearchIndex downloads and parses the search index from the docs site. +func DownloadSearchIndex(ctx context.Context, url string) (*SearchIndex, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return nil, fmt.Errorf("creating request: %w", err) + } + + resp, err := http.DefaultClient.Do(req) + if err != nil { + return nil, fmt.Errorf("fetching search index: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("unexpected status code: %d", resp.StatusCode) + } + + body, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("reading response body: %w", err) + } + + var index SearchIndex + if err := json.Unmarshal(body, &index); err != nil { + return nil, fmt.Errorf("parsing search index: %w", err) + } + + return &index, nil +} + +// ProcessSearchIndex processes the search index into pages, filtering out unwanted content. +func ProcessSearchIndex(index *SearchIndex, baseURL string) []Page { + // Group documents by page (location without fragment) + pageMap := make(map[string]*Page) + pageOrder := make([]string, 0) + + for _, doc := range index.Docs { + // Skip empty documents + if doc.Text == "" && doc.Title == "" { + continue + } + + // Extract page path (without fragment) + pagePath := getPagePath(doc.Location) + + // Skip tag index pages (they're just lists of links) + if isTagIndexPage(pagePath, doc.Location) { + continue + } + + // Get or create page + page, exists := pageMap[pagePath] + if !exists { + page = &Page{ + Path: pagePath, + URL: buildURL(baseURL, pagePath), + Sections: make([]Section, 0), + } + pageMap[pagePath] = page + pageOrder = append(pageOrder, pagePath) + } + + // Set page title from the root document (no fragment) + if !strings.Contains(doc.Location, "#") || page.Title == "" { + if doc.Title != "" { + page.Title = doc.Title + } + } + + // Collect tags + if len(doc.Tags) > 0 { + page.Tags = mergeTags(page.Tags, doc.Tags) + } + + // Add section content + if doc.Text != "" { + page.Sections = append(page.Sections, Section{ + Title: doc.Title, + Text: stripHTML(doc.Text), + }) + } + } + + // Convert map to slice in original order + pages := make([]Page, 0, len(pageOrder)) + for _, path := range pageOrder { + if page := pageMap[path]; page != nil && len(page.Sections) > 0 { + pages = append(pages, *page) + } + } + + return pages +} + +// ChunkPages converts pages into chunks suitable for embedding. +func ChunkPages(pages []Page, maxChars int) []Chunk { + chunks := make([]Chunk, 0) + + for _, page := range pages { + pageChunks := chunkPage(page, maxChars) + chunks = append(chunks, pageChunks...) + } + + return chunks +} + +// chunkPage splits a single page into chunks. +func chunkPage(page Page, maxChars int) []Chunk { + chunks := make([]Chunk, 0) + + // Build the header that will be included in each chunk + header := buildChunkHeader(page.Title, page.Tags) + headerLen := len(header) + + // Calculate available space for content + availableChars := maxChars - headerLen + if availableChars < 100 { + // If header is too long, use a smaller header + header = fmt.Sprintf("Title: %s\n\n", page.Title) + headerLen = len(header) + availableChars = maxChars - headerLen + } + + // Combine all section content + var contentBuilder strings.Builder + for i, section := range page.Sections { + if i > 0 { + contentBuilder.WriteString("\n\n") + } + if section.Title != "" && section.Title != page.Title { + contentBuilder.WriteString("## ") + contentBuilder.WriteString(section.Title) + contentBuilder.WriteString("\n") + } + contentBuilder.WriteString(section.Text) + } + fullContent := contentBuilder.String() + + // If content fits in one chunk, return it + if len(fullContent) <= availableChars { + chunks = append(chunks, Chunk{ + Title: page.Title, + URL: page.URL, + Content: header + fullContent, + }) + return chunks + } + + // Split content into multiple chunks + contentChunks := splitContent(fullContent, availableChars) + for _, content := range contentChunks { + chunks = append(chunks, Chunk{ + Title: page.Title, + URL: page.URL, + Content: header + content, + }) + } + + return chunks +} + +// buildChunkHeader creates the header for a chunk including title and tags. +func buildChunkHeader(title string, tags []string) string { + var builder strings.Builder + builder.WriteString("Title: ") + builder.WriteString(title) + builder.WriteString("\n") + + if len(tags) > 0 { + builder.WriteString("Tags: ") + builder.WriteString(strings.Join(tags, ", ")) + builder.WriteString("\n") + } + + builder.WriteString("\n") + return builder.String() +} + +// splitContent splits content into chunks at natural boundaries. +func splitContent(content string, maxChars int) []string { + if len(content) <= maxChars { + return []string{content} + } + + chunks := make([]string, 0) + remaining := content + + for len(remaining) > 0 { + if len(remaining) <= maxChars { + chunks = append(chunks, strings.TrimSpace(remaining)) + break + } + + // Find a good split point + splitPoint := findSplitPoint(remaining, maxChars) + chunk := strings.TrimSpace(remaining[:splitPoint]) + if chunk != "" { + chunks = append(chunks, chunk) + } + remaining = remaining[splitPoint:] + } + + return chunks +} + +// findSplitPoint finds a good point to split content, preferring natural boundaries. +func findSplitPoint(content string, maxChars int) int { + if len(content) <= maxChars { + return len(content) + } + + // Look for section boundaries (## headers) first + sectionPattern := regexp.MustCompile(`\n## `) + if loc := sectionPattern.FindStringIndex(content[maxChars/2 : maxChars]); loc != nil { + return maxChars/2 + loc[0] + } + + // Look for paragraph boundaries + for i := maxChars - 1; i > maxChars/2; i-- { + if i < len(content) && content[i] == '\n' && i+1 < len(content) && content[i+1] == '\n' { + return i + 1 + } + } + + // Look for sentence boundaries + for i := maxChars - 1; i > maxChars/2; i-- { + if i < len(content) && (content[i] == '.' || content[i] == '!' || content[i] == '?') { + if i+1 < len(content) && (content[i+1] == ' ' || content[i+1] == '\n') { + return i + 1 + } + } + } + + // Look for word boundaries + for i := maxChars - 1; i > maxChars/2; i-- { + if i < len(content) && content[i] == ' ' { + return i + 1 + } + } + + // Fall back to hard cut + return maxChars +} + +// getPagePath extracts the page path from a location (removes fragment). +func getPagePath(location string) string { + if idx := strings.Index(location, "#"); idx != -1 { + return location[:idx] + } + return location +} + +// isTagIndexPage checks if a document is a tag index page that should be skipped. +func isTagIndexPage(pagePath, location string) bool { + // Skip the main tags page and individual tag pages + if pagePath == "tags/" || strings.HasPrefix(pagePath, "tags/") { + return true + } + // Skip tag fragments on other pages + if strings.Contains(location, "#tag:") { + return true + } + return false +} + +// buildURL constructs the full URL for a page. +func buildURL(baseURL, path string) string { + if path == "" { + return baseURL + "/" + } + return baseURL + "/" + path +} + +// mergeTags merges two tag slices, removing duplicates. +func mergeTags(existing, new []string) []string { + seen := make(map[string]bool) + for _, t := range existing { + seen[t] = true + } + + result := append([]string{}, existing...) + for _, t := range new { + if !seen[t] { + seen[t] = true + result = append(result, t) + } + } + return result +} + +// HTML tag patterns for stripping +var ( + htmlTagPattern = regexp.MustCompile(`<[^>]*>`) + multiSpacePattern = regexp.MustCompile(`[ \t]+`) + multiNewlinePattern = regexp.MustCompile(`\n{3,}`) +) + +// stripHTML removes HTML tags and decodes HTML entities. +func stripHTML(text string) string { + // Remove HTML tags + text = htmlTagPattern.ReplaceAllString(text, "") + + // Decode HTML entities + text = html.UnescapeString(text) + + // Normalize whitespace + text = multiSpacePattern.ReplaceAllString(text, " ") + text = multiNewlinePattern.ReplaceAllString(text, "\n\n") + + return strings.TrimSpace(text) +} diff --git a/internal/cmd/gen_rag_index/process_test.go b/internal/cmd/gen_rag_index/process_test.go new file mode 100644 index 000000000..d2127f26b --- /dev/null +++ b/internal/cmd/gen_rag_index/process_test.go @@ -0,0 +1,403 @@ +package gen_rag_index + +import ( + "testing" +) + +func TestStripHTML(t *testing.T) { + tests := []struct { + name string + input string + expected string + }{ + { + name: "simple tags", + input: "

Hello world

", + expected: "Hello world", + }, + { + name: "nested tags", + input: "
  • Item 1
  • Item 2
", + expected: "Item 1Item 2", + }, + { + name: "html entities", + input: "Hello & goodbye <test>", + expected: "Hello & goodbye ", + }, + { + name: "mixed content", + input: "

This is code and "quotes"

", + expected: "This is code and \"quotes\"", + }, + { + name: "multiple spaces", + input: "Hello world test", + expected: "Hello world test", + }, + { + name: "empty string", + input: "", + expected: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := stripHTML(tt.input) + if result != tt.expected { + t.Errorf("stripHTML(%q) = %q, want %q", tt.input, result, tt.expected) + } + }) + } +} + +func TestGetPagePath(t *testing.T) { + tests := []struct { + name string + location string + expected string + }{ + { + name: "no fragment", + location: "auth/entra-id/", + expected: "auth/entra-id/", + }, + { + name: "with fragment", + location: "auth/entra-id/#spec", + expected: "auth/entra-id/", + }, + { + name: "empty string", + location: "", + expected: "", + }, + { + name: "only fragment", + location: "#welcome", + expected: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := getPagePath(tt.location) + if result != tt.expected { + t.Errorf("getPagePath(%q) = %q, want %q", tt.location, result, tt.expected) + } + }) + } +} + +func TestIsTagIndexPage(t *testing.T) { + tests := []struct { + name string + pagePath string + location string + expected bool + }{ + { + name: "tags main page", + pagePath: "tags/", + location: "tags/", + expected: true, + }, + { + name: "tag subpage", + pagePath: "tags/", + location: "tags/#tag:how-to", + expected: true, + }, + { + name: "regular page", + pagePath: "auth/", + location: "auth/", + expected: false, + }, + { + name: "regular page with fragment", + pagePath: "auth/", + location: "auth/#logging-in-users", + expected: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := isTagIndexPage(tt.pagePath, tt.location) + if result != tt.expected { + t.Errorf("isTagIndexPage(%q, %q) = %v, want %v", tt.pagePath, tt.location, result, tt.expected) + } + }) + } +} + +func TestBuildURL(t *testing.T) { + tests := []struct { + name string + baseURL string + path string + expected string + }{ + { + name: "simple path", + baseURL: "https://docs.example.cloud.nais.io", + path: "auth/", + expected: "https://docs.example.cloud.nais.io/auth/", + }, + { + name: "empty path", + baseURL: "https://docs.example.cloud.nais.io", + path: "", + expected: "https://docs.example.cloud.nais.io/", + }, + { + name: "nested path", + baseURL: "https://docs.example.cloud.nais.io", + path: "auth/entra-id/reference/", + expected: "https://docs.example.cloud.nais.io/auth/entra-id/reference/", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := buildURL(tt.baseURL, tt.path) + if result != tt.expected { + t.Errorf("buildURL(%q, %q) = %q, want %q", tt.baseURL, tt.path, result, tt.expected) + } + }) + } +} + +func TestMergeTags(t *testing.T) { + tests := []struct { + name string + existing []string + new []string + expected []string + }{ + { + name: "no overlap", + existing: []string{"a", "b"}, + new: []string{"c", "d"}, + expected: []string{"a", "b", "c", "d"}, + }, + { + name: "with overlap", + existing: []string{"a", "b"}, + new: []string{"b", "c"}, + expected: []string{"a", "b", "c"}, + }, + { + name: "empty existing", + existing: []string{}, + new: []string{"a", "b"}, + expected: []string{"a", "b"}, + }, + { + name: "empty new", + existing: []string{"a", "b"}, + new: []string{}, + expected: []string{"a", "b"}, + }, + { + name: "both empty", + existing: []string{}, + new: []string{}, + expected: []string{}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := mergeTags(tt.existing, tt.new) + if len(result) != len(tt.expected) { + t.Errorf("mergeTags(%v, %v) = %v, want %v", tt.existing, tt.new, result, tt.expected) + return + } + for i, v := range result { + if v != tt.expected[i] { + t.Errorf("mergeTags(%v, %v) = %v, want %v", tt.existing, tt.new, result, tt.expected) + return + } + } + }) + } +} + +func TestBuildChunkHeader(t *testing.T) { + tests := []struct { + name string + title string + tags []string + expected string + }{ + { + name: "with tags", + title: "My Page", + tags: []string{"auth", "how-to"}, + expected: "Title: My Page\nTags: auth, how-to\n\n", + }, + { + name: "without tags", + title: "My Page", + tags: nil, + expected: "Title: My Page\n\n", + }, + { + name: "empty tags", + title: "My Page", + tags: []string{}, + expected: "Title: My Page\n\n", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := buildChunkHeader(tt.title, tt.tags) + if result != tt.expected { + t.Errorf("buildChunkHeader(%q, %v) = %q, want %q", tt.title, tt.tags, result, tt.expected) + } + }) + } +} + +func TestProcessSearchIndex(t *testing.T) { + index := &SearchIndex{ + Docs: []SearchDoc{ + { + Location: "auth/", + Title: "Authentication", + Text: "

Auth overview

", + Tags: []string{"auth", "explanation"}, + }, + { + Location: "auth/#logging-in", + Title: "Logging in", + Text: "

How to log in

", + Tags: []string{"auth", "explanation"}, + }, + { + Location: "tags/", + Title: "Tags", + Text: "", + }, + { + Location: "tags/#tag:how-to", + Title: "how-to", + Text: "
  • Item
", + }, + }, + } + + pages := ProcessSearchIndex(index, "https://docs.test.cloud.nais.io") + + // Should have 1 page (auth/), tags should be skipped + if len(pages) != 1 { + t.Errorf("ProcessSearchIndex() returned %d pages, want 1", len(pages)) + return + } + + page := pages[0] + if page.Path != "auth/" { + t.Errorf("page.Path = %q, want %q", page.Path, "auth/") + } + if page.Title != "Authentication" { + t.Errorf("page.Title = %q, want %q", page.Title, "Authentication") + } + if len(page.Tags) != 2 { + t.Errorf("page.Tags = %v, want 2 tags", page.Tags) + } + if len(page.Sections) != 2 { + t.Errorf("page.Sections = %d, want 2", len(page.Sections)) + } + if page.URL != "https://docs.test.cloud.nais.io/auth/" { + t.Errorf("page.URL = %q, want %q", page.URL, "https://docs.test.cloud.nais.io/auth/") + } +} + +func TestChunkPages(t *testing.T) { + pages := []Page{ + { + Path: "short/", + Title: "Short Page", + URL: "https://example.com/short/", + Tags: []string{"test"}, + Sections: []Section{ + {Title: "Short Page", Text: "This is short content."}, + }, + }, + } + + chunks := ChunkPages(pages, 1500) + + if len(chunks) != 1 { + t.Errorf("ChunkPages() returned %d chunks, want 1", len(chunks)) + return + } + + chunk := chunks[0] + if chunk.Title != "Short Page" { + t.Errorf("chunk.Title = %q, want %q", chunk.Title, "Short Page") + } + if chunk.URL != "https://example.com/short/" { + t.Errorf("chunk.URL = %q, want %q", chunk.URL, "https://example.com/short/") + } + + // Should contain header with title and tags + if !contains(chunk.Content, "Title: Short Page") { + t.Errorf("chunk.Content should contain title header") + } + if !contains(chunk.Content, "Tags: test") { + t.Errorf("chunk.Content should contain tags header") + } + if !contains(chunk.Content, "This is short content.") { + t.Errorf("chunk.Content should contain the actual content") + } +} + +func TestChunkPagesLongContent(t *testing.T) { + // Create content that's longer than chunk size + longContent := "" + for i := 0; i < 100; i++ { + longContent += "This is a sentence that adds length to the content. " + } + + pages := []Page{ + { + Path: "long/", + Title: "Long Page", + URL: "https://example.com/long/", + Sections: []Section{ + {Title: "Long Page", Text: longContent}, + }, + }, + } + + chunks := ChunkPages(pages, 500) + + if len(chunks) < 2 { + t.Errorf("ChunkPages() should split long content into multiple chunks, got %d", len(chunks)) + } + + // All chunks should have the URL + for _, chunk := range chunks { + if chunk.URL != "https://example.com/long/" { + t.Errorf("all chunks should have the same URL") + } + } +} + +func contains(s, substr string) bool { + return len(s) >= len(substr) && (s == substr || len(s) > 0 && containsHelper(s, substr)) +} + +func containsHelper(s, substr string) bool { + for i := 0; i <= len(s)-len(substr); i++ { + if s[i:i+len(substr)] == substr { + return true + } + } + return false +} diff --git a/internal/cmd/gen_rag_index/writer.go b/internal/cmd/gen_rag_index/writer.go new file mode 100644 index 000000000..4e67fa97c --- /dev/null +++ b/internal/cmd/gen_rag_index/writer.go @@ -0,0 +1,180 @@ +package gen_rag_index + +import ( + "context" + "database/sql" + "encoding/binary" + "fmt" + "math" + "os" + "path/filepath" + + _ "github.com/duckdb/duckdb-go/v2" + "github.com/sirupsen/logrus" +) + +// Writer handles writing chunks and embeddings to a DuckDB file. +type Writer struct { + db *sql.DB + log logrus.FieldLogger + finalPath string + tempPath string +} + +// NewWriter creates a new DuckDB writer. +func NewWriter(outputPath string, log logrus.FieldLogger) (*Writer, error) { + // Ensure the output directory exists + dir := filepath.Dir(outputPath) + if dir != "" && dir != "." { + if err := os.MkdirAll(dir, 0o755); err != nil { + return nil, fmt.Errorf("creating output directory: %w", err) + } + } + + // Use a temporary file for writing + tempPath := outputPath + ".tmp" + + // Remove existing temp file if it exists + if _, err := os.Stat(tempPath); err == nil { + if err := os.Remove(tempPath); err != nil { + return nil, fmt.Errorf("removing existing temp file: %w", err) + } + } + + // Open DuckDB on temp path + db, err := sql.Open("duckdb", tempPath) + if err != nil { + return nil, fmt.Errorf("opening DuckDB: %w", err) + } + + // Create the schema + if err := createSchema(db); err != nil { + db.Close() + return nil, fmt.Errorf("creating schema: %w", err) + } + + log.WithField("temp_path", tempPath).Info("initialized DuckDB writer") + + return &Writer{ + db: db, + log: log, + finalPath: outputPath, + tempPath: tempPath, + }, nil +} + +// createSchema creates the rag_documents table. +func createSchema(db *sql.DB) error { + _, err := db.Exec(` + CREATE TABLE rag_documents ( + title TEXT NOT NULL, + url TEXT NOT NULL, + content TEXT NOT NULL, + embedding BLOB NOT NULL + ) + `) + return err +} + +// ChunkWithEmbedding represents a chunk with its embedding vector. +type ChunkWithEmbedding struct { + Chunk Chunk + Embedding []float32 +} + +// WriteChunks writes chunks with embeddings to the database. +func (w *Writer) WriteChunks(ctx context.Context, chunks []ChunkWithEmbedding) error { + tx, err := w.db.BeginTx(ctx, nil) + if err != nil { + return fmt.Errorf("beginning transaction: %w", err) + } + defer tx.Rollback() + + stmt, err := tx.PrepareContext(ctx, ` + INSERT INTO rag_documents (title, url, content, embedding) + VALUES (?, ?, ?, ?) + `) + if err != nil { + return fmt.Errorf("preparing statement: %w", err) + } + defer stmt.Close() + + for _, chunk := range chunks { + embeddingBytes := encodeEmbedding(chunk.Embedding) + _, err := stmt.ExecContext(ctx, chunk.Chunk.Title, chunk.Chunk.URL, chunk.Chunk.Content, embeddingBytes) + if err != nil { + return fmt.Errorf("inserting chunk %q: %w", chunk.Chunk.Title, err) + } + } + + if err := tx.Commit(); err != nil { + return fmt.Errorf("committing transaction: %w", err) + } + + w.log.WithField("count", len(chunks)).Info("wrote chunks to database") + return nil +} + +// CreateIndex creates an index on the url column for faster lookups. +func (w *Writer) CreateIndex(ctx context.Context) error { + _, err := w.db.ExecContext(ctx, `CREATE INDEX idx_url ON rag_documents(url)`) + if err != nil { + return fmt.Errorf("creating index: %w", err) + } + w.log.Info("created index on url column") + return nil +} + +// GetStats returns statistics about the written data. +func (w *Writer) GetStats(ctx context.Context) (Stats, error) { + var stats Stats + + row := w.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM rag_documents`) + if err := row.Scan(&stats.TotalChunks); err != nil { + return stats, fmt.Errorf("counting chunks: %w", err) + } + + row = w.db.QueryRowContext(ctx, `SELECT COUNT(DISTINCT url) FROM rag_documents`) + if err := row.Scan(&stats.UniqueURLs); err != nil { + return stats, fmt.Errorf("counting unique URLs: %w", err) + } + + row = w.db.QueryRowContext(ctx, `SELECT AVG(LENGTH(content)) FROM rag_documents`) + if err := row.Scan(&stats.AvgContentLength); err != nil { + return stats, fmt.Errorf("calculating avg content length: %w", err) + } + + return stats, nil +} + +// Stats contains statistics about the generated index. +type Stats struct { + TotalChunks int + UniqueURLs int + AvgContentLength float64 +} + +// Close closes the database connection and atomically moves the file to the final destination. +func (w *Writer) Close() error { + if err := w.db.Close(); err != nil { + return fmt.Errorf("closing database: %w", err) + } + + if err := os.Rename(w.tempPath, w.finalPath); err != nil { + return fmt.Errorf("renaming temp file to final path: %w", err) + } + + w.log.WithField("final_path", w.finalPath).Info("atomically wrote index file") + return nil +} + +// encodeEmbedding converts a float32 slice to bytes for storage. +// Uses little-endian format to match the searcher's parseEmbedding function. +func encodeEmbedding(embedding []float32) []byte { + buf := make([]byte, len(embedding)*4) + for i, v := range embedding { + bits := math.Float32bits(v) + binary.LittleEndian.PutUint32(buf[i*4:], bits) + } + return buf +} diff --git a/internal/database/migrations/0060_agent_conversations.sql b/internal/database/migrations/0060_agent_conversations.sql new file mode 100644 index 000000000..c7a617343 --- /dev/null +++ b/internal/database/migrations/0060_agent_conversations.sql @@ -0,0 +1,29 @@ +-- +goose Up +-- Conversation threads for the agent chat +CREATE TABLE agent_conversations ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + title TEXT NOT NULL DEFAULT 'New conversation', + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX idx_agent_conversations_user_id ON agent_conversations(user_id); +CREATE INDEX idx_agent_conversations_user_created ON agent_conversations(user_id, created_at); + +-- Individual messages within a conversation +CREATE TABLE agent_messages ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + conversation_id UUID NOT NULL REFERENCES agent_conversations(id) ON DELETE CASCADE, + role TEXT NOT NULL CHECK (role IN ('user', 'assistant', 'tool')), + content TEXT NOT NULL, + sources JSONB, + blocks JSONB, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX idx_agent_messages_conversation_id ON agent_messages(conversation_id); + +-- +goose Down +DROP TABLE IF EXISTS agent_messages; +DROP TABLE IF EXISTS agent_conversations; diff --git a/internal/grpc/grpc.go b/internal/grpc/grpc.go index 99026e63b..2c22f7cc5 100644 --- a/internal/grpc/grpc.go +++ b/internal/grpc/grpc.go @@ -18,9 +18,16 @@ import ( "google.golang.org/grpc" ) -func Run(ctx context.Context, listenAddress string, pool *pgxpool.Pool, log logrus.FieldLogger) error { - log.Info("GRPC serving on ", listenAddress) - lis, err := net.Listen("tcp", listenAddress) +// Config holds configuration for the gRPC server. +type Config struct { + ListenAddress string + Pool *pgxpool.Pool + Log logrus.FieldLogger +} + +func Run(ctx context.Context, cfg *Config) error { + cfg.Log.Info("GRPC serving on ", cfg.ListenAddress) + lis, err := net.Listen("tcp", cfg.ListenAddress) if err != nil { return fmt.Errorf("failed to listen: %w", err) } @@ -30,10 +37,10 @@ func Run(ctx context.Context, listenAddress string, pool *pgxpool.Pool, log logr } s := grpc.NewServer(opts...) - protoapi.RegisterTeamsServer(s, grpcteam.NewServer(pool)) - protoapi.RegisterUsersServer(s, grpcuser.NewServer(pool)) - protoapi.RegisterReconcilersServer(s, grpcreconciler.NewServer(pool)) - protoapi.RegisterDeploymentsServer(s, grpcdeployment.NewServer(pool)) + protoapi.RegisterTeamsServer(s, grpcteam.NewServer(cfg.Pool)) + protoapi.RegisterUsersServer(s, grpcuser.NewServer(cfg.Pool)) + protoapi.RegisterReconcilersServer(s, grpcreconciler.NewServer(cfg.Pool)) + protoapi.RegisterDeploymentsServer(s, grpcdeployment.NewServer(cfg.Pool)) g, ctx := errgroup.WithContext(ctx) g.Go(func() error { return s.Serve(lis) })