verbose with history
This commit is contained in:
32
main.py
32
main.py
@@ -2,26 +2,50 @@ import os
|
|||||||
import sys
|
import sys
|
||||||
from dotenv import load_dotenv
|
from dotenv import load_dotenv
|
||||||
from google import genai
|
from google import genai
|
||||||
|
from google.genai import types
|
||||||
|
|
||||||
|
messages = []
|
||||||
|
|
||||||
load_dotenv()
|
load_dotenv()
|
||||||
api_key = os.environ.get("GEMINI_API_KEY")
|
api_key = os.environ.get("GEMINI_API_KEY")
|
||||||
|
|
||||||
client = genai.Client(api_key=api_key)
|
client = genai.Client(api_key=api_key)
|
||||||
|
|
||||||
|
def add_message(user, message):
|
||||||
|
messages.append({"role": user, "parts": [{"text": message}]})
|
||||||
|
|
||||||
def main():
|
def main():
|
||||||
if len(sys.argv) < 2:
|
if len(sys.argv) < 2:
|
||||||
print("Error: No prompt provided.\nUsage: uv run main.py \"<your prompt here>\"")
|
print("Error: No prompt provided.\nUsage: uv run main.py \"<your prompt here>\" [--verbose]")
|
||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
|
|
||||||
user_prompt = sys.argv[1]
|
args = sys.argv[1:]
|
||||||
|
|
||||||
|
verbose = False
|
||||||
|
if "--verbose" in args:
|
||||||
|
verbose = True
|
||||||
|
args.remove("--verbose")
|
||||||
|
|
||||||
|
user_prompt = " ".join(args)
|
||||||
|
|
||||||
|
if verbose:
|
||||||
|
print(f'User prompt: "{user_prompt}"')
|
||||||
|
|
||||||
|
add_message("user", user_prompt)
|
||||||
|
|
||||||
response = client.models.generate_content(
|
response = client.models.generate_content(
|
||||||
model="gemini-2.0-flash-001",
|
model="gemini-2.0-flash-001",
|
||||||
contents=user_prompt
|
contents=messages
|
||||||
)
|
)
|
||||||
|
|
||||||
print(response.text)
|
reply_text = response.text
|
||||||
|
|
||||||
|
# Add AI response to history
|
||||||
|
add_message("model", reply_text)
|
||||||
|
|
||||||
|
print(reply_text)
|
||||||
|
|
||||||
|
if verbose:
|
||||||
print(f"Prompt tokens: {response.usage_metadata.prompt_token_count}")
|
print(f"Prompt tokens: {response.usage_metadata.prompt_token_count}")
|
||||||
print(f"Response tokens: {response.usage_metadata.candidates_token_count}")
|
print(f"Response tokens: {response.usage_metadata.candidates_token_count}")
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user