Inference plugin: Add Gemini model adapter (elastic#191292)

## Summary Add the `gemini` model adapter for the `inference` plugin. Had to perform minor changes on the associated connector Also update the codeowner files to add the `@elastic/appex-ai-infra` team as (one of the) owner of the genAI connectors --------- Co-authored-by: kibanamachine <[email protected]>
juliaElastic · Aug 28, 2024 · 02a3992 · 02a3992
1 parent fc93e43
commit 02a3992
Show file tree

Hide file tree

Showing 19 changed files with 1,261 additions and 56 deletions.
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
@@ -1557,18 +1557,18 @@ x-pack/test/security_solution_cypress/cypress/tasks/expandable_flyout  @elastic/
 
 ## Generative AI owner connectors
 # OpenAI
-/x-pack/plugins/stack_connectors/public/connector_types/openai @elastic/security-generative-ai @elastic/obs-ai-assistant
-/x-pack/plugins/stack_connectors/server/connector_types/openai @elastic/security-generative-ai @elastic/obs-ai-assistant
-/x-pack/plugins/stack_connectors/common/openai @elastic/security-generative-ai @elastic/obs-ai-assistant
+/x-pack/plugins/stack_connectors/public/connector_types/openai @elastic/security-generative-ai @elastic/obs-ai-assistant @elastic/appex-ai-infra
+/x-pack/plugins/stack_connectors/server/connector_types/openai @elastic/security-generative-ai @elastic/obs-ai-assistant @elastic/appex-ai-infra
+/x-pack/plugins/stack_connectors/common/openai @elastic/security-generative-ai @elastic/obs-ai-assistant @elastic/appex-ai-infra
 # Bedrock
-/x-pack/plugins/stack_connectors/public/connector_types/bedrock @elastic/security-generative-ai @elastic/obs-ai-assistant
-/x-pack/plugins/stack_connectors/server/connector_types/bedrock @elastic/security-generative-ai @elastic/obs-ai-assistant
-/x-pack/plugins/stack_connectors/common/bedrock @elastic/security-generative-ai @elastic/obs-ai-assistant
+/x-pack/plugins/stack_connectors/public/connector_types/bedrock @elastic/security-generative-ai @elastic/obs-ai-assistant @elastic/appex-ai-infra
+/x-pack/plugins/stack_connectors/server/connector_types/bedrock @elastic/security-generative-ai @elastic/obs-ai-assistant @elastic/appex-ai-infra
+/x-pack/plugins/stack_connectors/common/bedrock @elastic/security-generative-ai @elastic/obs-ai-assistant @elastic/appex-ai-infra
 
 # Gemini
-/x-pack/plugins/stack_connectors/public/connector_types/gemini @elastic/security-generative-ai @elastic/obs-ai-assistant
-/x-pack/plugins/stack_connectors/server/connector_types/gemini @elastic/security-generative-ai @elastic/obs-ai-assistant
-/x-pack/plugins/stack_connectors/common/gemini @elastic/security-generative-ai @elastic/obs-ai-assistant
+/x-pack/plugins/stack_connectors/public/connector_types/gemini @elastic/security-generative-ai @elastic/obs-ai-assistant @elastic/appex-ai-infra
+/x-pack/plugins/stack_connectors/server/connector_types/gemini @elastic/security-generative-ai @elastic/obs-ai-assistant @elastic/appex-ai-infra
+/x-pack/plugins/stack_connectors/common/gemini @elastic/security-generative-ai @elastic/obs-ai-assistant @elastic/appex-ai-infra
 
 ## Defend Workflows owner connectors
 /x-pack/plugins/stack_connectors/public/connector_types/sentinelone @elastic/security-defend-workflows

diff --git a/x-pack/plugins/actions/server/integration_tests/__snapshots__/connector_types.test.ts.snap b/x-pack/plugins/actions/server/integration_tests/__snapshots__/connector_types.test.ts.snap
diff --git a/x-pack/plugins/inference/common/chat_complete/tool_schema.ts b/x-pack/plugins/inference/common/chat_complete/tool_schema.ts
@@ -5,15 +5,15 @@
  * 2.0.
  */
 
-import { Required, ValuesType, UnionToIntersection } from 'utility-types';
+import { Required, ValuesType } from 'utility-types';
 
 interface ToolSchemaFragmentBase {
   description?: string;
 }
 
 interface ToolSchemaTypeObject extends ToolSchemaFragmentBase {
   type: 'object';
-  properties: Record<string, ToolSchemaFragment>;
+  properties: Record<string, ToolSchemaType>;
   required?: string[] | readonly string[];
 }
 
@@ -35,28 +35,18 @@ interface ToolSchemaTypeNumber extends ToolSchemaFragmentBase {
   enum?: string[] | readonly string[];
 }
 
-interface ToolSchemaAnyOf extends ToolSchemaFragmentBase {
-  anyOf: ToolSchemaType[];
-}
-
-interface ToolSchemaAllOf extends ToolSchemaFragmentBase {
-  allOf: ToolSchemaType[];
-}
-
 interface ToolSchemaTypeArray extends ToolSchemaFragmentBase {
   type: 'array';
   items: Exclude<ToolSchemaType, ToolSchemaTypeArray>;
 }
 
-type ToolSchemaType =
+export type ToolSchemaType =
   | ToolSchemaTypeObject
   | ToolSchemaTypeString
   | ToolSchemaTypeBoolean
   | ToolSchemaTypeNumber
   | ToolSchemaTypeArray;
 
-type ToolSchemaFragment = ToolSchemaType | ToolSchemaAnyOf | ToolSchemaAllOf;
-
 type FromToolSchemaObject<TToolSchemaObject extends ToolSchemaTypeObject> = Required<
   {
     [key in keyof TToolSchemaObject['properties']]?: FromToolSchema<
@@ -79,17 +69,9 @@ type FromToolSchemaString<TToolSchemaString extends ToolSchemaTypeString> =
     ? ValuesType<TToolSchemaString['enum']>
     : string;
 
-type FromToolSchemaAnyOf<TToolSchemaAnyOf extends ToolSchemaAnyOf> = FromToolSchema<
-  ValuesType<TToolSchemaAnyOf['anyOf']>
->;
-
-type FromToolSchemaAllOf<TToolSchemaAllOf extends ToolSchemaAllOf> = UnionToIntersection<
-  FromToolSchema<ValuesType<TToolSchemaAllOf['allOf']>>
->;
-
 export type ToolSchema = ToolSchemaTypeObject;
 
-export type FromToolSchema<TToolSchema extends ToolSchemaFragment> =
+export type FromToolSchema<TToolSchema extends ToolSchemaType> =
   TToolSchema extends ToolSchemaTypeObject
     ? FromToolSchemaObject<TToolSchema>
     : TToolSchema extends ToolSchemaTypeArray
@@ -100,8 +82,4 @@ export type FromToolSchema<TToolSchema extends ToolSchemaFragment> =
     ? number
     : TToolSchema extends ToolSchemaTypeString
     ? FromToolSchemaString<TToolSchema>
-    : TToolSchema extends ToolSchemaAnyOf
-    ? FromToolSchemaAnyOf<TToolSchema>
-    : TToolSchema extends ToolSchemaAllOf
-    ? FromToolSchemaAllOf<TToolSchema>
     : never;
diff --git a/x-pack/plugins/inference/server/chat_complete/adapters/gemini/gemini_adapter.test.mocks.ts b/x-pack/plugins/inference/server/chat_complete/adapters/gemini/gemini_adapter.test.mocks.ts
@@ -0,0 +1,16 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export const processVertexStreamMock = jest.fn();
+
+jest.doMock('./process_vertex_stream', () => {
+  const actual = jest.requireActual('./process_vertex_stream');
+  return {
+    ...actual,
+    processVertexStream: processVertexStreamMock,
+  };
+});