File size: 632 Bytes
c9f26e8
 
 
 
 
 
 
 
 
 
 
4f5f090
 
 
 
 
 
c9f26e8
 
4f5f090
 
 
 
c9f26e8
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
import toml
from google import genai

def load_prompt(args):
    with open(args.prompt_tmpl_path, 'r') as f:
        prompts = toml.load(f)

    return prompts

def setup_gemini_client(args):
    if args.vertexai:
        client = genai.client.AsyncClient(
            genai.client.ApiClient(
                vertexai=args.vertexai, 
                project=args.vertexai_project, 
                location=args.vertexai_location
            )
        )
    else:
        client = genai.client.AsyncClient(
            genai.client.ApiClient(
                api_key=args.ai_studio_api_key
            )
        )

    return client