in internal/enricher/enricher.go [675:728]
func (mc *MetadataCollector) generateTableDescriptionWithGemini(ctx context.Context, metadata *TableMetadata, schemaContext string) (string, error) {
if mc.GeminiAPIKey == "" {
return "", nil
}
if mc.AdditionalContext == "" {
return "", nil
}
prompt := fmt.Sprintf(`
Your task is to generate a brief and concise description for a database table based on the provided context.
The context might be irrelevant, so you need to firstly read through the context and decide if there is any relevant information for the target table.
********** Knowledge Context **********
%s
********** End Knowledge Context **********
**Instructions:**
- Response starting with your analysis. Then output the description in between <result></result> tags.
- Be concise and informative, no more than 50 words.
- Important: Only provide a description for the table if there is information related to this table in additional context. Otherwise in all other cases, return empty <result></result> tags.
The target table is:
Table: %s
Now start your response. Remember, only give description when the knowledge context provides useful information about the table.
`, mc.AdditionalContext, metadata.Table)
client, err := genai.NewClient(ctx, option.WithAPIKey(mc.GeminiAPIKey))
if err != nil {
return "", fmt.Errorf("failed to create Gemini client: %w", err)
}
defer client.Close()
model_name := mc.Model
if model_name == "" {
model_name = "gemini-1.5-pro-002"
}
model := client.GenerativeModel(model_name)
model.SetTemperature(0.4)
model.SetMaxOutputTokens(500)
model.SetTopP(0.8)
model.SetTopK(40)
resp, err := model.GenerateContent(ctx, genai.Text(prompt))
if err != nil {
return "", fmt.Errorf("Gemini API call failed: %w", err)
}
description, err := extractTextFromResponse(resp)
if err != nil {
return "", err
}
return description, nil
}