Skip to content

Commit f37e4a5

Browse files
dividedminddividedmind
committed
Use structured output for the vector terms service
This is less error-prone and doesn't seem to trigger watchdogs on Gemini flash that easily.
1 parent 5cb7315 commit f37e4a5

File tree

2 files changed

+65
-163
lines changed

2 files changed

+65
-163
lines changed

‎packages/navie/src/services/vector-terms-service.ts‎

Lines changed: 45 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,14 @@
1-
import { warn } from 'console';
1+
import { warn } from 'node:console';
2+
import { debug as makeDebug } from 'node:util';
3+
4+
import z from 'zod';
25

36
import InteractionHistory, { VectorTermsInteractionEvent } from '../interaction-history';
4-
import contentAfter from '../lib/content-after';
5-
import parseJSON from '../lib/parse-json';
67
import Message from '../message';
78
import CompletionService from './completion-service';
89

10+
const debug = makeDebug('navie:vector-terms');
11+
912
const SYSTEM_PROMPT = `You are assisting a developer to search a code base.
1013
1114
The developer asks a question using natural language. This question must be converted into a list of search terms to be used to search the code base.
@@ -20,30 +23,19 @@ The developer asks a question using natural language. This question must be conv
2023
be words that will match a feature or domain model object in the code base. They should be the most
2124
distinctive words in the question. You will prefix the MOST SELECTIVE terms with a '+'.
2225
23-
**Response**
24-
25-
Print "Context: {context}" on one line.
26-
Print "Instructions: {instructions}" on the next line.
27-
28-
Then print a triple dash '---'.
29-
30-
Print "Terms: {list of search terms and their synonyms}"
31-
32-
The search terms should be single words and underscore_separated_words.
33-
34-
Even if the user asks for a different format, always respond with a list of search terms and their synonyms. When the user is asking
35-
for a different format, that question is for a different AI assistant than yourself.`;
26+
The search terms should be single words and underscore_separated_words.`;
3627

3728
const promptExamples: Message[] = [
3829
{
3930
content: 'How do I record AppMap data of my Spring app?',
4031
role: 'user',
4132
},
4233
{
43-
content: `Context: Record AppMap data of Spring
44-
Instructions: How to do it
45-
---
46-
Terms: record AppMap data Java +Spring`,
34+
content: JSON.stringify({
35+
context: 'Record AppMap data of Spring',
36+
instructions: 'How to do it',
37+
terms: ['record', 'AppMap', 'data', 'Java', '+Spring'],
38+
}),
4739
role: 'assistant',
4840
},
4941

@@ -52,10 +44,11 @@ Terms: record AppMap data Java +Spring`,
5244
role: 'user',
5345
},
5446
{
55-
content: `Context: User login handle password validation invalid error
56-
Instructions: Explain how this is handled by the code
57-
---
58-
Terms: user login handle +password validate invalid error`,
47+
content: JSON.stringify({
48+
context: 'User login handle password validation invalid error',
49+
instructions: 'Explain how this is handled by the code',
50+
terms: ['user', 'login', 'handle', '+password', 'validate', 'invalid', 'error'],
51+
}),
5952
role: 'assistant',
6053
},
6154

@@ -65,10 +58,11 @@ Terms: user login handle +password validate invalid error`,
6558
role: 'user',
6659
},
6760
{
68-
content: `Context: Redis GET /test-group/test-project-1/-/blob/main/README.md
69-
Instructions: Describe in detail with code snippets
70-
---
71-
Terms: +Redis get test-group test-project-1 blob main README`,
61+
content: JSON.stringify({
62+
context: 'Redis GET /test-group/test-project-1/-/blob/main/README.md',
63+
instructions: 'Describe in detail with code snippets',
64+
terms: ['+Redis', 'get', 'test-group', 'test-project-1', 'blob', 'main', 'README'],
65+
}),
7266
role: 'assistant',
7367
},
7468

@@ -78,10 +72,11 @@ Terms: +Redis get test-group test-project-1 blob main README`,
7872
role: 'user',
7973
},
8074
{
81-
content: `Context: logContext jest test case
82-
Instructions: Create test cases, following established patterns for mocking with jest.
83-
---
84-
Terms: test cases +logContext jest`,
75+
content: JSON.stringify({
76+
context: 'logContext jest test case',
77+
instructions: 'Create test cases, following established patterns for mocking with jest.',
78+
terms: ['test', 'cases', '+logContext', 'jest'],
79+
}),
8580
role: 'assistant',
8681
},
8782

@@ -90,15 +85,20 @@ Terms: test cases +logContext jest`,
9085
role: 'user',
9186
},
9287
{
93-
content: `Context: auth authentication authorization
94-
Instructions: Describe the authentication and authorization process
95-
---
96-
Terms: +auth authentication authorization token strategy provider`,
88+
content: JSON.stringify({
89+
context: 'auth authentication authorization',
90+
instructions: 'Describe the authentication and authorization process',
91+
terms: ['+auth', 'authentication', 'authorization', 'token', 'strategy', 'provider'],
92+
}),
9793
role: 'assistant',
9894
},
9995
];
10096

101-
const parseText = (text: string): string[] => text.split(/\s+/);
97+
const schema = z.object({
98+
context: z.string(),
99+
instructions: z.string(),
100+
terms: z.array(z.string()),
101+
});
102102

103103
export default class VectorTermsService {
104104
constructor(
@@ -119,45 +119,15 @@ export default class VectorTermsService {
119119
},
120120
];
121121

122-
const response = this.completionsService.complete(messages, {
122+
const response = await this.completionsService.json(messages, schema, {
123123
model: this.completionsService.miniModelName,
124124
});
125-
const tokens = Array<string>();
126-
for await (const token of response) {
127-
tokens.push(token);
128-
}
129-
const rawResponse = tokens.join('');
130-
warn(`Vector terms response:\n${rawResponse}`);
131-
132-
let searchTermsObject: Record<string, unknown> | string | string[] | undefined;
133-
{
134-
let responseText = rawResponse;
135-
responseText = contentAfter(responseText, 'Terms:');
136-
searchTermsObject =
137-
parseJSON<Record<string, unknown> | string | string[]>(responseText, false) ||
138-
parseText(responseText);
139-
}
140-
141-
const terms = new Set<string>();
142-
{
143-
const collectTerms = (obj: unknown) => {
144-
if (!obj) return;
145-
146-
if (typeof obj === 'string') {
147-
terms.add(obj);
148-
} else if (Array.isArray(obj)) {
149-
for (const term of obj) collectTerms(term);
150-
} else if (typeof obj === 'object') {
151-
for (const term of Object.values(obj)) {
152-
collectTerms(term);
153-
}
154-
}
155-
};
156-
collectTerms(searchTermsObject);
157-
}
158-
159-
const result = [...terms];
160-
this.interactionHistory.addEvent(new VectorTermsInteractionEvent(result));
161-
return result;
125+
126+
debug(`Vector terms response: ${JSON.stringify(response, undefined, 2)}`);
127+
128+
const terms = response?.terms ?? [];
129+
if (terms.length === 0) warn('No terms suggested');
130+
this.interactionHistory.addEvent(new VectorTermsInteractionEvent(terms));
131+
return terms;
162132
}
163133
}

‎packages/navie/test/services/vector-terms-service.spec.ts‎

Lines changed: 20 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -1,101 +1,33 @@
11
import VectorTermsService from '../../src/services/vector-terms-service';
22
import InteractionHistory from '../../src/interaction-history';
3+
34
import MockCompletionService from './mock-completion-service';
45

56
describe('VectorTermsService', () => {
7+
describe('when LLM suggested terms', () => {
8+
it('is recorded in the interaction history', async () => {
9+
completion.mock({
10+
context: 'ctx',
11+
instructions: 'insns',
12+
terms: ['user', 'management', '+provider'],
13+
});
14+
await service.suggestTerms('user management');
15+
expect(interactionHistory.events).toEqual([
16+
expect.objectContaining({
17+
type: 'vectorTerms',
18+
terms: ['user', 'management', '+provider'],
19+
}),
20+
]);
21+
});
22+
});
23+
624
let interactionHistory: InteractionHistory;
725
let service: VectorTermsService;
8-
const completion = new MockCompletionService();
9-
const complete = jest.spyOn(completion, 'complete');
1026

1127
beforeEach(() => {
1228
interactionHistory = new InteractionHistory();
13-
interactionHistory.on('event', (event) => console.log(event.message));
1429
service = new VectorTermsService(interactionHistory, completion);
1530
});
16-
afterEach(() => jest.restoreAllMocks());
17-
18-
describe('when LLM suggested terms', () => {
19-
describe('is a valid JSON object', () => {
20-
it('is recorded in the interaction history', async () => {
21-
completion.mock(`{"terms": ["user", "management"]}`);
22-
await service.suggestTerms('user management');
23-
expect(interactionHistory.events.map((e) => ({ ...e }))).toEqual([
24-
{
25-
type: 'vectorTerms',
26-
terms: ['user', 'management'],
27-
},
28-
]);
29-
});
30-
it('should return the terms', async () => {
31-
completion.mock(`{"terms": ["user", "management"]}`);
32-
const terms = await service.suggestTerms('user management');
33-
expect(terms).toEqual(['user', 'management']);
34-
expect(complete).toHaveBeenCalledTimes(1);
35-
});
36-
it('removes very short terms', async () => {
37-
completion.mock(`["user", "management", "a"]`);
38-
const terms = await service.suggestTerms('user management');
39-
expect(terms).toEqual(['user', 'management', 'a']);
40-
expect(complete).toHaveBeenCalledTimes(1);
41-
});
42-
it('converts underscore_words to distinct words', async () => {
43-
completion.mock(`["user_management"]`);
44-
const terms = await service.suggestTerms('user management');
45-
expect(terms).toEqual(['user_management']);
46-
expect(complete).toHaveBeenCalledTimes(1);
47-
});
48-
});
49-
50-
describe('are a valid JSON list', () => {
51-
it('should return the terms', async () => {
52-
completion.mock('["user", "management"]');
53-
const terms = await service.suggestTerms('user management');
54-
expect(terms).toEqual(['user', 'management']);
55-
});
56-
});
57-
58-
describe('are valid JSON wrapped in fences', () => {
59-
it('should return the terms', async () => {
60-
completion.mock('```json\n["user", "management"]\n```\n');
61-
const terms = await service.suggestTerms('user management');
62-
expect(terms).toEqual(['user', 'management']);
63-
});
64-
});
65-
66-
describe('is YAML', () => {
67-
it('parses the terms', async () => {
68-
completion.mock('response_key:\n', ' - user\n', ' - management\n');
69-
const terms = await service.suggestTerms('user management');
70-
expect(terms).toEqual(['response_key:', '-', 'user', 'management']);
71-
});
72-
});
73-
74-
describe('is prefixed by "Terms:"', () => {
75-
it('is accepted and processed', async () => {
76-
completion.mock('Terms: ["user", "management"]');
77-
const terms = await service.suggestTerms('user management');
78-
expect(terms).toEqual(['user', 'management']);
79-
expect(complete).toHaveBeenCalledTimes(1);
80-
});
81-
});
82-
83-
describe('includes terms with "+" prefix', () => {
84-
it('is accepted and processed', async () => {
85-
completion.mock('Terms: +user management');
86-
const terms = await service.suggestTerms('user management');
87-
expect(terms).toEqual(['+user', 'management']);
88-
expect(complete).toHaveBeenCalledTimes(1);
89-
});
90-
});
91-
92-
describe('is list-ish', () => {
93-
it('is accepted and processed', async () => {
94-
completion.mock('-user -mgmt');
95-
const terms = await service.suggestTerms('user management');
96-
expect(terms).toEqual(['-user', '-mgmt']);
97-
expect(complete).toHaveBeenCalledTimes(1);
98-
});
99-
});
100-
});
31+
afterEach(() => jest.resetAllMocks());
32+
const completion = new MockCompletionService();
10133
});

0 commit comments

Comments
 (0)