kshitijthakkar commited on
Commit
68c1485
·
1 Parent(s): 807fe76

fix: Use correct GeminiClient method in generate_prompt_template

Browse files

- Changed from non-existent generate_content() to model.generate_content_async()
- Added proper generation_config dict
- Matches pattern used in generate_synthetic_dataset tool
- Fixes AttributeError when calling the tool

Files changed (1) hide show
  1. mcp_tools.py +10 -3
mcp_tools.py CHANGED
@@ -2020,11 +2020,18 @@ I have a base {template_name} prompt template and need to customize it for a spe
2020
  Start your response with the YAML content immediately."""
2021
 
2022
  # Call Gemini to customize the template
2023
- customized_template = await gemini_client.generate_content(
 
 
 
 
 
 
 
2024
  customization_prompt,
2025
- temperature=0.3, # Lower temperature for more consistent formatting
2026
- max_output_tokens=4096
2027
  )
 
2028
 
2029
  # Clean up the response (remove any markdown formatting if present)
2030
  customized_template = customized_template.strip()
 
2020
  Start your response with the YAML content immediately."""
2021
 
2022
  # Call Gemini to customize the template
2023
+ generation_config = {
2024
+ "temperature": 0.3, # Lower temperature for more consistent formatting
2025
+ "top_p": 0.95,
2026
+ "top_k": 40,
2027
+ "max_output_tokens": 4096,
2028
+ }
2029
+
2030
+ response = await gemini_client.model.generate_content_async(
2031
  customization_prompt,
2032
+ generation_config=generation_config
 
2033
  )
2034
+ customized_template = response.text
2035
 
2036
  # Clean up the response (remove any markdown formatting if present)
2037
  customized_template = customized_template.strip()