Skip to content

Commit 8b8982d

Browse files
hlshenhlshen
committed
initial orcas
1 parent e28281f commit 8b8982d

4 files changed

Lines changed: 395 additions & 113 deletions

File tree

‎firebase-vscode/src/data-connect/execution/execution.ts‎

Lines changed: 39 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,10 @@ import { DataConnectError, toSerializedError } from "../../../common/error";
2626
import { InstanceType } from "../code-lens-provider";
2727
import { DATA_CONNECT_EVENT_NAME, AnalyticsLogger } from "../../analytics";
2828
import { EmulatorsController } from "../../core/emulators";
29-
import { getConnectorGQLText, insertQueryAt } from "../file-utils";
29+
import { getConnectorGQLText, insertQueryAt, findGqlFiles } from "../file-utils";
30+
import * as fs from "fs";
31+
import * as path from "path";
32+
import { dataConnectConfigs } from "../config";
3033
import * as gif from "../../../../src/gemini/fdcExperience";
3134
import { ensureGIFApiTos } from "../../../../src/dataconnect/ensureApis";
3235
import { configstore } from "../../../../src/configstore";
@@ -282,17 +285,41 @@ export function registerExecution(
282285
return;
283286
}
284287
try {
285-
const schema = await dataConnectService.schema();
286-
// TODO: Update to SQL Connect once backend agent is updated
288+
const configs = dataConnectConfigs.value?.tryReadValue;
289+
const serviceConfig = configs?.findEnclosingServiceForPath(arg.document.fileName);
290+
291+
const schemas: any[] = [];
292+
293+
if (serviceConfig) {
294+
const mainSchemaDir = path.join(serviceConfig.path, serviceConfig.mainSchemaDir);
295+
const secondaryDirs = serviceConfig.secondarySchemaDirs.map(dir => path.join(serviceConfig.path, dir));
296+
297+
const schemaFiles: string[] = [];
298+
schemaFiles.push(...(await findGqlFiles(mainSchemaDir)));
299+
for (const dir of secondaryDirs) {
300+
schemaFiles.push(...(await findGqlFiles(dir)));
301+
}
302+
303+
const files = await Promise.all(schemaFiles.map(async (file) => {
304+
const content = fs.readFileSync(file, "utf-8");
305+
return {
306+
path: path.basename(file),
307+
content: content
308+
};
309+
}));
310+
311+
if (files.length > 0) {
312+
schemas.push({
313+
source: {
314+
files: files
315+
}
316+
});
317+
}
318+
}
319+
287320
const prompt = `Generate a Data Connect operation to match this description: ${arg.description}
288-
${arg.existingQuery ? `\n\nRefine this existing operation:\n${arg.existingQuery}` : ""}
289-
${
290-
schema
291-
? `\n\nUse the Data Connect Schema:\n\`\`\`graphql
292-
${schema}
293-
\`\`\``
294-
: ""
295-
}`;
321+
${arg.existingQuery ? `\n\nRefine this existing operation:\n${arg.existingQuery}` : ""}`;
322+
296323
const serviceName = await dataConnectService.servicePath(
297324
arg.document.fileName,
298325
);
@@ -305,6 +332,7 @@ ${schema}
305332
prompt,
306333
serviceName,
307334
arg.projectId,
335+
schemas.length > 0 ? schemas : undefined
308336
);
309337
await insertQueryAt(
310338
arg.document.uri,

‎src/gemini/fdcExperience.spec.ts‎

Lines changed: 182 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -1,65 +1,194 @@
11
import { expect } from "chai";
2-
import { extractCodeBlock } from "./fdcExperience";
3-
4-
describe("extractCodeBlock", () => {
5-
it("should extract a basic GraphQL query block", () => {
6-
const text =
7-
'Here is a GraphQL query:\n```graphql\nquery GetUser { user(id: "1") { name email } }\n```\nThanks!';
8-
const expected = 'query GetUser { user(id: "1") { name email } }';
9-
expect(extractCodeBlock(text)).to.eq(expected);
2+
import { extractCodeBlock, generateSchema, generateOperation } from "./fdcExperience";
3+
import * as nock from "nock";
4+
5+
describe("fdcExperience", () => {
6+
beforeEach(() => {
7+
nock.cleanAll();
8+
});
9+
10+
after(() => {
11+
nock.cleanAll();
1012
});
1113

12-
it("should extract a multi-line GraphQL mutation block", () => {
13-
const text = `
14-
Some preamble.
15-
\`\`\`graphql
16-
mutation CreatePost($title: String!, $content: String!) {
17-
createPost(title: $title, content: $content) {
18-
id
19-
title
14+
describe("extractCodeBlock", () => {
15+
it("should extract a basic GraphQL query block", () => {
16+
const text =
17+
'Here is a GraphQL query:\n```graphql\nquery GetUser { user(id: "1") { name email } }\n```\nThanks!';
18+
const expected = 'query GetUser { user(id: "1") { name email } }';
19+
expect(extractCodeBlock(text)).to.eq(expected);
20+
});
21+
22+
it("should extract a multi-line GraphQL mutation block", () => {
23+
const text = `
24+
Some preamble.
25+
\`\`\`graphql
26+
mutation CreatePost($title: String!, $content: String!) {
27+
createPost(title: $title, content: $content) {
28+
id
29+
title
30+
}
31+
}
32+
\`\`\`
33+
Followed by some description.
34+
`;
35+
const expected = `mutation CreatePost($title: String!, $content: String!) {
36+
createPost(title: $title, content: $content) {
37+
id
38+
title
39+
}
40+
}`;
41+
expect(extractCodeBlock(text)).to.eq(expected);
42+
});
43+
44+
it("should extract a GraphQL fragment block", () => {
45+
const text = "```graphql\nfragment UserFields on User { id name }\n```";
46+
const expected = "fragment UserFields on User { id name }";
47+
expect(extractCodeBlock(text)).to.eq(expected);
48+
});
49+
50+
it("should extract an empty GraphQL code block", () => {
51+
const text = "```graphql\n\n```";
52+
const expected = "";
53+
expect(extractCodeBlock(text)).to.eq(expected);
54+
});
55+
56+
it("should extract a GraphQL schema definition block", () => {
57+
const text = `
58+
\`\`\`graphql
59+
type Query {
60+
hello: String
2061
}
21-
}
22-
\`\`\`
23-
Followed by some description.
24-
`;
25-
const expected = `mutation CreatePost($title: String!, $content: String!) {
26-
createPost(title: $title, content: $content) {
27-
id
28-
title
62+
schema {
63+
query: Query
2964
}
30-
}`;
31-
expect(extractCodeBlock(text)).to.eq(expected);
65+
\`\`\`
66+
`;
67+
const expected = `type Query {
68+
hello: String
69+
}
70+
schema {
71+
query: Query
72+
}`;
73+
expect(extractCodeBlock(text)).to.eq(expected);
74+
});
3275
});
3376

34-
it("should extract a GraphQL fragment block", () => {
35-
const text = "```graphql\nfragment UserFields on User { id name }\n```";
36-
const expected = "fragment UserFields on User { id name }";
37-
expect(extractCodeBlock(text)).to.eq(expected);
38-
});
77+
describe("generateSchema", () => {
78+
it("should make a streaming POST request and return parsed code", async () => {
79+
const prompt = "Create a blog";
80+
const project = "my-project";
81+
const location = "us-central1";
82+
const responseObj = {
83+
part: {
84+
codeChunk: {
85+
code: "type User { id: String }",
86+
languageCode: "graphql"
87+
}
88+
}
89+
};
3990

40-
it("should extract an empty GraphQL code block", () => {
41-
const text = "```graphql\n\n```";
42-
const expected = "";
43-
expect(extractCodeBlock(text)).to.eq(expected);
91+
nock("https://autopush-firebasedataconnect.sandbox.googleapis.com")
92+
.post(`/v1/projects/${project}/locations/${location}/services/-:generateSchema`, {
93+
name: `projects/${project}/locations/${location}/services/-`,
94+
prompt
95+
})
96+
.reply(200, JSON.stringify(responseObj));
97+
98+
const result = await generateSchema(prompt, project, location);
99+
expect(result).to.equal("type User { id: String }");
100+
expect(nock.isDone()).to.be.true;
101+
});
102+
103+
it("should call onStatus callback when status updates are received", async () => {
104+
const prompt = "Create a blog";
105+
const project = "my-project";
106+
const location = "us-central1";
107+
const statusObj = {
108+
status: {
109+
state: "ANALYZING_SCHEMA",
110+
message: "Analyzing schema..."
111+
}
112+
};
113+
const responseObj = {
114+
part: {
115+
codeChunk: {
116+
code: "type User { id: String }",
117+
languageCode: "graphql"
118+
}
119+
}
120+
};
121+
122+
let statusCalledWith: any = null;
123+
const onStatus = (status: any) => {
124+
statusCalledWith = status;
125+
};
126+
127+
nock("https://autopush-firebasedataconnect.sandbox.googleapis.com")
128+
.post(`/v1/projects/${project}/locations/${location}/services/-:generateSchema`, {
129+
name: `projects/${project}/locations/${location}/services/-`,
130+
prompt
131+
})
132+
.reply(200, JSON.stringify(statusObj) + "\n" + JSON.stringify(responseObj));
133+
134+
const result = await generateSchema(prompt, project, location, onStatus);
135+
expect(result).to.equal("type User { id: String }");
136+
expect(statusCalledWith).to.deep.equal(statusObj.status);
137+
expect(nock.isDone()).to.be.true;
138+
});
44139
});
45140

46-
it("should extract a GraphQL schema definition block", () => {
47-
const text = `
48-
\`\`\`graphql
49-
type Query {
50-
hello: String
51-
}
52-
schema {
53-
query: Query
54-
}
55-
\`\`\`
56-
`;
57-
const expected = `type Query {
58-
hello: String
59-
}
60-
schema {
61-
query: Query
62-
}`;
63-
expect(extractCodeBlock(text)).to.eq(expected);
141+
describe("generateOperation", () => {
142+
it("should make a streaming POST request and return parsed code", async () => {
143+
const prompt = "Get users";
144+
const service = "projects/my-project/locations/us-central1/services/my-service";
145+
const project = "my-project";
146+
const responseObj = {
147+
part: {
148+
codeChunk: {
149+
code: "query GetUsers { users { id } }",
150+
languageCode: "graphql"
151+
}
152+
}
153+
};
154+
155+
nock("https://autopush-firebasedataconnect.sandbox.googleapis.com")
156+
.post(`/v1/projects/my-project/locations/us-central1/services/my-service:generateQuery`, {
157+
name: `projects/my-project/locations/us-central1/services/my-service`,
158+
prompt
159+
})
160+
.reply(200, JSON.stringify(responseObj));
161+
162+
const result = await generateOperation(prompt, service, project);
163+
expect(result).to.equal("query GetUsers { users { id } }");
164+
expect(nock.isDone()).to.be.true;
165+
});
166+
167+
it("should use '-' as serviceId when schemas are provided", async () => {
168+
const prompt = "Get users";
169+
const service = "projects/my-project/locations/us-central1/services/my-service";
170+
const project = "my-project";
171+
const schemas = [{ source: { files: [{ path: "schema.gql", content: "type User { id: String }" }] } }];
172+
const responseObj = {
173+
part: {
174+
codeChunk: {
175+
code: "query GetUsers { users { id } }",
176+
languageCode: "graphql"
177+
}
178+
}
179+
};
180+
181+
nock("https://autopush-firebasedataconnect.sandbox.googleapis.com")
182+
.post(`/v1/projects/my-project/locations/us-central1/services/-:generateQuery`, {
183+
name: `projects/my-project/locations/us-central1/services/-`,
184+
prompt,
185+
schemas
186+
})
187+
.reply(200, JSON.stringify(responseObj));
188+
189+
const result = await generateOperation(prompt, service, project, schemas);
190+
expect(result).to.equal("query GetUsers { users { id } }");
191+
expect(nock.isDone()).to.be.true;
192+
});
64193
});
65194
});

0 commit comments

Comments
 (0)