File size: 2,662 Bytes
920dfd0 |
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 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
import os
import json
from google import genai
from google.genai import types
from dotenv import load_dotenv
from prompts import *
try:
from dotenv import load_dotenv
load_dotenv()
except ImportError:
pass
def format_search_results(results):
"""Format Search Results"""
formatted_output = ""
for i in range(1, 4):
formatted_output += f"### {i}. {results[f'keyword{i}']}\n"
formatted_output += f"{results[f'summary{i}']}"
formatted_output += "\n"
return formatted_output
def grounding_with_google_search(stt_data, content_type = "국민의힘"):
"""
Extract Keywords and Perform Google Search
"""
# Create a client for Google GenAI
client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
# Define a sample schema for search results
sample_schema = {
"keyword1": "Core Keyword 1",
"summary1": "Summarize Search Result about Core Keyword 1",
"keyword2": "Core Keyword 2",
"summary2": "Summarize Search Result about Core Keyword 2",
"keyword3": "Core Keyword 3",
"summary3": "Summarize Search Result about Core Keyword 3"
}
# Select the appropriate system prompt based on content_type
if content_type == "Agents for Amazon Bedrock":
system_prompt = BEDROCK_SEARCH_PROMPT
elif content_type == "Bundesliga Fan Experience":
system_prompt = BUNDESLIGA_SEARCH_PROMPT
elif content_type == "AWS_2024_recap":
system_prompt = AWS_SEARCH_PROMPT
# Format the system prompt with the sample schema
system_prompt = system_prompt.format(sample_schema=sample_schema)
# Prepare the human message with the input script
human_message = f"""
## Input Script
{stt_data}
"""
# Generate content using the Google GenAI client
response = client.models.generate_content(
model="gemini-2.0-flash-001",
contents=human_message,
config=types.GenerateContentConfig(
system_instruction=system_prompt,
response_mime_type="application/json",
tools=[
types.Tool(
google_search=types.GoogleSearchRetrieval(
dynamic_retrieval_config = types.DynamicRetrievalConfig(
mode=types.DynamicRetrievalConfigMode.MODE_UNSPECIFIED,
dynamic_threshold=0.0
)
)
)
]
)
)
# Parse the response text and format the search results
text = response.text
results = json.loads(text)
return format_search_results(results) |