@@ -43,10 +43,7 @@ def _query_gemini(self, user_query, functions):
43
43
}
44
44
45
45
# NOTE: To run the gemini model, you need to provide your own GCP project ID, which can be found in the GCP console.
46
- if self .model_name == "gemini-1.5-pro-preview-0409" :
47
- API_URL = "https://us-central1-aiplatform.googleapis.com/v1beta1/projects/{YOUR_GCP_PROJECT_ID_HERE}/locations/us-central1/publishers/google/models/gemini-1.5-pro-preview-0409:generateContent"
48
- else :
49
- API_URL = "https://us-central1-aiplatform.googleapis.com/v1beta1/projects/{YOUR_GCP_PROJECT_ID_HERE}/locations/us-central1/publishers/google/models/gemini-1.0-pro:generateContent"
46
+ API_URL = "https://us-central1-aiplatform.googleapis.com/v1beta1/projects/{YOUR_GCP_PROJECT_ID_HERE}/locations/us-central1/publishers/google/models/" + self .model_name + ":generateContent"
50
47
headers = {
51
48
"Authorization" : "Bearer " + token ,
52
49
"Content-Type" : "application/json" ,
@@ -65,22 +62,19 @@ def _query_gemini(self, user_query, functions):
65
62
"output_tokens" : 0 ,
66
63
"latency" : latency ,
67
64
}
68
- contents = result ["candidates" ][0 ]["content" ]["parts" ][0 ]
69
- if "functionCall" in contents :
70
- if (
71
- "name" in contents ["functionCall" ]
72
- and "args" in contents ["functionCall" ]
73
- ):
74
- result = {
75
- contents ["functionCall" ]["name" ]: json .dumps (
76
- contents ["functionCall" ]["args" ]
77
- )
78
- }
79
-
65
+ parts = []
66
+ for part in result ["candidates" ][0 ]["content" ]["parts" ]:
67
+ if "functionCall" in part :
68
+ if (
69
+ "name" in part ["functionCall" ]
70
+ and "args" in part ["functionCall" ]
71
+ ):
72
+ parts .append ({part ["functionCall" ]["name" ]: json .dumps (part ["functionCall" ]["args" ])})
73
+ else :
74
+ parts .append ("Parsing error: " + json .dumps (part ["functionCall" ]))
80
75
else :
81
- result = "Parsing error: " + json .dumps (contents ["functionCall" ])
82
- else :
83
- result = contents ["text" ]
76
+ parts .append (part ["text" ])
77
+ result = parts
84
78
metatdata = {}
85
79
metatdata ["input_tokens" ] = json .loads (response .content )["usageMetadata" ][
86
80
"promptTokenCount"
0 commit comments