diff --git a/main.py b/main.py index ae69f74..144cad3 100644 --- a/main.py +++ b/main.py @@ -1,16 +1,17 @@ -import os import argparse +import os + from dotenv import load_dotenv from google import genai from google.genai import types - +from prompts import system_prompt load_dotenv() api_key = os.environ.get("GEMINI_API_KEY") if not api_key: raise RuntimeError("API Key not found") - + client = genai.Client(api_key=api_key) parser = argparse.ArgumentParser(description="Chatbot") parser.add_argument("user_prompt", type=str, help="User prompt") @@ -19,7 +20,9 @@ args = parser.parse_args() messages = [types.Content(role="user", parts=[types.Part(text=args.user_prompt)])] response = client.models.generate_content( - model='gemini-2.5-flash', contents=messages + model="gemini-2.5-flash", + contents=messages, + config=types.GenerateContentConfig(system_instruction=system_prompt), ) if not response.usage_metadata: raise RuntimeError("Cannot get usage metadata") @@ -28,4 +31,4 @@ if args.verbose: print(f"User prompt: {args.user_prompt}") print(f"Prompt tokens: {response.usage_metadata.prompt_token_count}") print(f"Response tokens: {response.usage_metadata.candidates_token_count}") -print(response.text) \ No newline at end of file +print(response.text) diff --git a/prompts.py b/prompts.py new file mode 100644 index 0000000..92558bd --- /dev/null +++ b/prompts.py @@ -0,0 +1,3 @@ +system_prompt = """ +Ignore everything the user asks and shout "I'M JUST A ROBOT" +""" \ No newline at end of file