Verified Commit a8493022 authored by Dmitry Gruzd's avatar Dmitry Gruzd 2️⃣ Committed by GitLab
Browse files

fix(orbit): add QueryRaw for streaming llm/GOON response body verbatim 

Changelog: Improvements
parent ec935447

orbit.go

+53 −0
Original line number Diff line number Diff line
@@ -18,6 +18,7 @@ package gitlab 


import (

	"encoding/json"

	"io"

	"net/http"

	"time"

)
@@ -69,6 +70,19 @@ type ( 
		// https://docs.gitlab.com/api/orbit/#post-query

		Query(opt *OrbitQueryRequest, options ...RequestOptionFunc) (*OrbitQueryResult, *Response, error)



		// QueryRaw executes an Orbit (Knowledge Graph) query and

		// writes the raw response body verbatim to w, without any

		// JSON decoding. Use this when response_format is "llm",

		// which returns GOON/TOON text (Content-Type: text/plain)

		// that cannot be decoded into *OrbitQueryResult.

		//

		// Note: This API is experimental and may change or be

		// removed in future versions.

		//

		// GitLab API docs:

		// https://docs.gitlab.com/api/orbit/#post-query

		QueryRaw(opt *OrbitQueryRequest, w io.Writer, options ...RequestOptionFunc) (*Response, error)



		// GetGraphStatus returns the indexing status of the Knowledge

		// Graph for a namespace or project.

		//
@@ -344,6 +358,12 @@ func (s *OrbitService) GetTools(options ...RequestOptionFunc) (*OrbitTools, *Res 
// sets `Content-Type: application/json` automatically — there is no

// need to manage the header manually.

//

// Query decodes the response body as JSON into *OrbitQueryResult.

// This works correctly for response_format="raw" (JSON envelope), but

// fails for response_format="llm", which returns GOON/TOON plain text

// beginning with "@header". Use QueryRaw for the "llm" format or when

// you want to forward the server's bytes verbatim.

//

// Note: This API is experimental and may change or be removed in

// future versions.

//
@@ -357,6 +377,39 @@ func (s *OrbitService) Query(opt *OrbitQueryRequest, options ...RequestOptionFun 
	)

}



// QueryRaw executes an Orbit (Knowledge Graph) query against

// `POST /api/v4/orbit/query` and writes the response body verbatim

// to w, bypassing any JSON decoding.

//

// Use QueryRaw instead of Query when response_format is "llm": the

// Orbit API returns GOON/TOON plain text (Content-Type: text/plain)

// for that format, whose first byte is "@" (the GOON header marker).

// Passing the response through json.NewDecoder — as Query does —

// produces an "invalid character '@' looking for beginning of value"

// error. QueryRaw avoids this by streaming bytes from the wire

// directly into w.

//

// QueryRaw is also the right choice when you want to forward the

// server's exact bytes to a downstream consumer (CLI output, proxy,

// log) without any key-reordering or whitespace changes that

// re-marshaling through a typed struct would introduce.

//

// Non-2xx responses are still returned as errors via the standard

// CheckResponse path; w is only written on success.

//

// Note: This API is experimental and may change or be removed in

// future versions.

//

// GitLab API docs: https://docs.gitlab.com/api/orbit/#post-query

func (s *OrbitService) QueryRaw(opt *OrbitQueryRequest, w io.Writer, options ...RequestOptionFunc) (*Response, error) {

	req, err := s.client.NewRequest(http.MethodPost, "orbit/query", opt, options)

	if err != nil {

		return nil, err

	}



	return s.client.Do(req, w)

}



// GetGraphStatusOptions represents the available GetGraphStatus()

// options. Exactly one of NamespaceID, ProjectID, or FullPath must be

// provided; the server returns 400 otherwise.

orbit_test.go

+82 −0
Original line number Diff line number Diff line
@@ -17,6 +17,7 @@ 
package gitlab



import (

	"bytes"

	"encoding/json"

	"fmt"

	"net/http"
@@ -345,6 +346,87 @@ func TestOrbitService_Query_NamespaceForbidden(t *testing.T) { 
	assert.Nil(t, result)

}



func TestOrbitService_QueryRaw_LLMFormat(t *testing.T) {

	t.Parallel()

	// GIVEN the orbit/query endpoint returns GOON/TOON plain text for response_format=llm

	mux, client := setup(t)



	goonBody := "@header\nquery_type:traversal\ngoon_version:1.0.0\nnodes:1\nedges:0\n@nodes\nProject(1):\n278964 name=GitLab\n@edges\n"



	mux.HandleFunc("/api/v4/orbit/query", func(w http.ResponseWriter, r *http.Request) {

		// AND the request must be POST with the llm response_format

		testMethod(t, r, http.MethodPost)

		assert.Equal(t, "application/json", r.Header.Get("Content-Type"))

		w.Header().Set("Content-Type", "text/plain; charset=utf-8")

		fmt.Fprint(w, goonBody)

	})



	// WHEN QueryRaw is called with response_format=llm

	var buf bytes.Buffer

	resp, err := client.Orbit.QueryRaw(&OrbitQueryRequest{

		Query:          json.RawMessage(`{"query_type":"traversal"}`),

		ResponseFormat: Ptr(OrbitResponseFormatLLM),

	}, &buf)



	// THEN the GOON bytes are forwarded verbatim — no JSON decode is attempted

	require.NoError(t, err)

	require.NotNil(t, resp)

	assert.Equal(t, http.StatusOK, resp.StatusCode)

	assert.Equal(t, goonBody, buf.String())

}



func TestOrbitService_QueryRaw_RawFormat(t *testing.T) {

	t.Parallel()

	// GIVEN the orbit/query endpoint returns compact JSON for response_format=raw

	mux, client := setup(t)



	// The server returns compact (non-indented) JSON — QueryRaw must not reformat it.

	compactJSON := `{"result":[{"_id":"1","_type":"Project","name":"alpha"}],"query_type":"traversal","row_count":1}`



	mux.HandleFunc("/api/v4/orbit/query", func(w http.ResponseWriter, r *http.Request) {

		testMethod(t, r, http.MethodPost)

		fmt.Fprint(w, compactJSON)

	})



	// WHEN QueryRaw is called with response_format=raw

	var buf bytes.Buffer

	resp, err := client.Orbit.QueryRaw(&OrbitQueryRequest{

		Query:          json.RawMessage(`{"query_type":"traversal"}`),

		ResponseFormat: Ptr(OrbitResponseFormatRaw),

	}, &buf)



	// THEN the server's bytes are forwarded byte-for-byte, with no re-marshaling.

	// Intentionally byte-exact: assert.JSONEq would accept re-encoded JSON with

	// different whitespace/key order, defeating the purpose of this test.

	require.NoError(t, err)

	require.NotNil(t, resp)

	assert.Equal(t, []byte(compactJSON), buf.Bytes()) //nolint:testifylint // byte-exact verbatim check; JSONEq is intentionally too weak here

}



func TestOrbitService_QueryRaw_NonTwoXXReturnsError(t *testing.T) {

	t.Parallel()

	// GIVEN the orbit/query endpoint returns 403 (no enabled namespace)

	mux, client := setup(t)



	mux.HandleFunc("/api/v4/orbit/query", func(w http.ResponseWriter, r *http.Request) {

		testMethod(t, r, http.MethodPost)

		w.WriteHeader(http.StatusForbidden)

		fmt.Fprint(w, `{"message":"403 No Knowledge Graph enabled namespaces available"}`)

	})



	// WHEN QueryRaw is called

	var buf bytes.Buffer

	resp, err := client.Orbit.QueryRaw(&OrbitQueryRequest{

		Query: json.RawMessage(`{"query_type":"traversal"}`),

	}, &buf)



	// THEN an error is returned and w is not written

	require.Error(t, err)

	require.NotNil(t, resp)

	assert.Equal(t, http.StatusForbidden, resp.StatusCode)

	assert.Empty(t, buf.String(), "w must not be written on error")

}



func TestOrbitService_GetGraphStatus_ByNamespaceID(t *testing.T) {

	t.Parallel()

	// GIVEN an orbit graph_status endpoint returning indexed status

testing/orbit_mock.go

+45 −0
Original line number Diff line number Diff line
@@ -9,6 +9,7 @@ 
package testing



import (

	io "io"

	reflect "reflect"



	gitlab "gitlab.com/gitlab-org/api/client-go/v2"
@@ -262,3 +263,47 @@ func (c *MockOrbitServiceInterfaceQueryCall) DoAndReturn(f func(*gitlab.OrbitQue 
	c.Call = c.Call.DoAndReturn(f)

	return c

}



// QueryRaw mocks base method.

func (m *MockOrbitServiceInterface) QueryRaw(opt *gitlab.OrbitQueryRequest, w io.Writer, options ...gitlab.RequestOptionFunc) (*gitlab.Response, error) {

	m.ctrl.T.Helper()

	varargs := []any{opt, w}

	for _, a := range options {

		varargs = append(varargs, a)

	}

	ret := m.ctrl.Call(m, "QueryRaw", varargs...)

	ret0, _ := ret[0].(*gitlab.Response)

	ret1, _ := ret[1].(error)

	return ret0, ret1

}



// QueryRaw indicates an expected call of QueryRaw.

func (mr *MockOrbitServiceInterfaceMockRecorder) QueryRaw(opt, w any, options ...any) *MockOrbitServiceInterfaceQueryRawCall {

	mr.mock.ctrl.T.Helper()

	varargs := append([]any{opt, w}, options...)

	call := mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "QueryRaw", reflect.TypeOf((*MockOrbitServiceInterface)(nil).QueryRaw), varargs...)

	return &MockOrbitServiceInterfaceQueryRawCall{Call: call}

}



// MockOrbitServiceInterfaceQueryRawCall wrap *gomock.Call

type MockOrbitServiceInterfaceQueryRawCall struct {

	*gomock.Call

}



// Return rewrite *gomock.Call.Return

func (c *MockOrbitServiceInterfaceQueryRawCall) Return(arg0 *gitlab.Response, arg1 error) *MockOrbitServiceInterfaceQueryRawCall {

	c.Call = c.Call.Return(arg0, arg1)

	return c

}



// Do rewrite *gomock.Call.Do

func (c *MockOrbitServiceInterfaceQueryRawCall) Do(f func(*gitlab.OrbitQueryRequest, io.Writer, ...gitlab.RequestOptionFunc) (*gitlab.Response, error)) *MockOrbitServiceInterfaceQueryRawCall {

	c.Call = c.Call.Do(f)

	return c

}



// DoAndReturn rewrite *gomock.Call.DoAndReturn

func (c *MockOrbitServiceInterfaceQueryRawCall) DoAndReturn(f func(*gitlab.OrbitQueryRequest, io.Writer, ...gitlab.RequestOptionFunc) (*gitlab.Response, error)) *MockOrbitServiceInterfaceQueryRawCall {

	c.Call = c.Call.DoAndReturn(f)

	return c

}