regend client

BoundaryML · Jul 16, 2024 · 27d35c8 · 27d35c8
1 parent 69d9ca9
commit 27d35c8
Show file tree

Hide file tree

Showing 13 changed files with 162 additions and 156 deletions.
diff --git a/docs/docs/snippets/clients/providers/vertex.mdx b/docs/docs/snippets/clients/providers/vertex.mdx
@@ -3,7 +3,15 @@ title: vertex-ai
 slug: docs/snippets/clients/providers/vertex
 ---
 
-The `vertex-ai` provider is used to interact with the Google Vertex AI services.
+The `vertex-ai` provider is used to interact with the Google Vertex AI services, specifically the following endpoints:
+
+`
+https://${LOCATION}-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/${LOCATION}/publishers/google/models/${MODEL_ID}:streamGenerateContent
+`
+
+`https://${LOCATION}-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/${LOCATION}/publishers/google/models/${MODEL_ID}:generateContent`
+
+
 
 Example: 
 ```baml BAML

diff --git a/engine/baml-runtime/src/internal/llm_client/primitive/mod.rs b/engine/baml-runtime/src/internal/llm_client/primitive/mod.rs
@@ -103,8 +103,8 @@ impl TryFrom<(&ClientProperty, &RuntimeContext)> for LLMPrimitiveProvider {
             "google-ai" => {
                 GoogleAIClient::dynamic_new(value, ctx).map(LLMPrimitiveProvider::Google)
             }
-            // "aws-bedrock" => aws::AwsClient::dynamic_new(value, ctx).map(LLMPrimitiveProvider::Aws),
             "vertex-ai" => VertexClient::dynamic_new(value, ctx).map(LLMPrimitiveProvider::Vertex),
+            // dynamic_new is not implemented for aws::AwsClient
             other => {
                 let options = [
                     "openai",

diff --git a/engine/baml-runtime/src/internal/llm_client/primitive/vertex/vertex_client.rs b/engine/baml-runtime/src/internal/llm_client/primitive/vertex/vertex_client.rs
@@ -428,7 +428,7 @@ impl RequestBuilder for VertexClient {
         };
 
         let credentials = self.properties.service_account_details.clone();
-        log::info!("Credentials: {:?}", credentials);
+
         let access_token = if let Some((key, value)) = &credentials {
             if key == "GOOGLE_TOKEN" {
                 value.clone()
@@ -647,8 +647,6 @@ fn convert_chat_prompt_to_body(
             .collect::<serde_json::Value>(),
     );
 
-    log::debug!("converted chat prompt to body: {:#?}", map);
-
     return map;
 }
 

diff --git a/engine/language_client_typescript/native.d.ts b/engine/language_client_typescript/native.d.ts
@@ -113,10 +113,11 @@ export interface BamlLogEvent {
   startTime: string
 }
 
-export function invoke_runtime_cli(params: Array<string>): void
+export declare function invoke_runtime_cli(params: Array<string>): void
 
 export interface LogEventMetadata {
   eventId: string
   parentId?: string
   rootEventId: string
 }
+
diff --git a/integ-tests/baml_src/clients.baml b/integ-tests/baml_src/clients.baml
@@ -89,8 +89,7 @@ client<llm> Gemini {
 }
 
 client<llm> Vertex {
-  provider vertex-ai
-
+  provider vertex-ai  
   options {
     model gemini-1.5-pro
     project_id anish-testing-426119

diff --git a/integ-tests/baml_src/test-files/functions/input/named-args/single/named-int.baml b/integ-tests/baml_src/test-files/functions/input/named-args/single/named-int.baml
@@ -1,6 +1,6 @@
 // test for int
 function TestFnNamedArgsSingleInt(myInt: int) -> string {
-  client Vertex
+  client GPT35
   prompt #"
     Return this value back to me: {{myInt}}
   "#

diff --git a/integ-tests/baml_src/test-files/providers/providers.baml b/integ-tests/baml_src/test-files/providers/providers.baml
@@ -5,7 +5,7 @@ function TestAnthropic(input: string) -> string {
   "#
 }
 
-function TestOpenAI(input: string) -> string {
+function PromptTestOpenAI(input: string) -> string {
   client GPT35
   prompt #"
     Write a nice haiku about {{ input }}

diff --git a/integ-tests/python/baml_client/client.py b/integ-tests/python/baml_client/client.py
@@ -992,6 +992,30 @@ async def PromptTestClaudeChatNoSystem(
       mdl = create_model("PromptTestClaudeChatNoSystemReturnType", inner=(str, ...))
       return coerce(mdl, raw.parsed())
 
+    async def PromptTestOpenAI(
+        self,
+        input: str,
+        baml_options: BamlCallOptions = {},
+    ) -> str:
+      __tb__ = baml_options.get("tb", None)
+      if __tb__ is not None:
+        tb = __tb__._tb
+      else:
+        tb = None
+      __cr__ = baml_options.get("client_registry", None)
+
+      raw = await self.__runtime.call_function(
+        "PromptTestOpenAI",
+        {
+          "input": input,
+        },
+        self.__ctx_manager.get(),
+        tb,
+        __cr__,
+      )
+      mdl = create_model("PromptTestOpenAIReturnType", inner=(str, ...))
+      return coerce(mdl, raw.parsed())
+
     async def PromptTestOpenAIChat(
         self,
         input: str,
@@ -1448,30 +1472,6 @@ async def TestOllama(
       mdl = create_model("TestOllamaReturnType", inner=(str, ...))
       return coerce(mdl, raw.parsed())
 
-    async def TestOpenAI(
-        self,
-        input: str,
-        baml_options: BamlCallOptions = {},
-    ) -> str:
-      __tb__ = baml_options.get("tb", None)
-      if __tb__ is not None:
-        tb = __tb__._tb
-      else:
-        tb = None
-      __cr__ = baml_options.get("client_registry", None)
-
-      raw = await self.__runtime.call_function(
-        "TestOpenAI",
-        {
-          "input": input,
-        },
-        self.__ctx_manager.get(),
-        tb,
-        __cr__,
-      )
-      mdl = create_model("TestOpenAIReturnType", inner=(str, ...))
-      return coerce(mdl, raw.parsed())
-
     async def TestOpenAILegacyProvider(
         self,
         input: str,
@@ -2893,6 +2893,39 @@ def PromptTestClaudeChatNoSystem(
         self.__ctx_manager.get(),
       )
 
+    def PromptTestOpenAI(
+        self,
+        input: str,
+        baml_options: BamlCallOptions = {},
+    ) -> baml_py.BamlStream[Optional[str], str]:
+      __tb__ = baml_options.get("tb", None)
+      if __tb__ is not None:
+        tb = __tb__._tb
+      else:
+        tb = None
+      __cr__ = baml_options.get("client_registry", None)
+
+      raw = self.__runtime.stream_function(
+        "PromptTestOpenAI",
+        {
+          "input": input,
+        },
+        None,
+        self.__ctx_manager.get(),
+        tb,
+        __cr__,
+      )
+
+      mdl = create_model("PromptTestOpenAIReturnType", inner=(str, ...))
+      partial_mdl = create_model("PromptTestOpenAIPartialReturnType", inner=(Optional[str], ...))
+
+      return baml_py.BamlStream[Optional[str], str](
+        raw,
+        lambda x: coerce(partial_mdl, x),
+        lambda x: coerce(mdl, x),
+        self.__ctx_manager.get(),
+      )
+
     def PromptTestOpenAIChat(
         self,
         input: str,
@@ -3520,39 +3553,6 @@ def TestOllama(
         self.__ctx_manager.get(),
       )
 
-    def TestOpenAI(
-        self,
-        input: str,
-        baml_options: BamlCallOptions = {},
-    ) -> baml_py.BamlStream[Optional[str], str]:
-      __tb__ = baml_options.get("tb", None)
-      if __tb__ is not None:
-        tb = __tb__._tb
-      else:
-        tb = None
-      __cr__ = baml_options.get("client_registry", None)
-
-      raw = self.__runtime.stream_function(
-        "TestOpenAI",
-        {
-          "input": input,
-        },
-        None,
-        self.__ctx_manager.get(),
-        tb,
-        __cr__,
-      )
-
-      mdl = create_model("TestOpenAIReturnType", inner=(str, ...))
-      partial_mdl = create_model("TestOpenAIPartialReturnType", inner=(Optional[str], ...))
-
-      return baml_py.BamlStream[Optional[str], str](
-        raw,
-        lambda x: coerce(partial_mdl, x),
-        lambda x: coerce(mdl, x),
-        self.__ctx_manager.get(),
-      )
-
     def TestOpenAILegacyProvider(
         self,
         input: str,