From 5398027ecb27ae87cb8bcf782210698f6f19387e Mon Sep 17 00:00:00 2001 From: Abijah Date: Wed, 3 Jun 2026 15:10:05 -0700 Subject: [PATCH] Working on the gemini CLI --- gemini/gemini.py | 161 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 161 insertions(+) create mode 100644 gemini/gemini.py diff --git a/gemini/gemini.py b/gemini/gemini.py new file mode 100644 index 0000000..eb3722c --- /dev/null +++ b/gemini/gemini.py @@ -0,0 +1,161 @@ +#!/usr/bin/env python3 +import argparse +import json +import os +import sys +import datetime +import google.generativeai as genai +from google.generativeai import caching + +def main(): + parser = argparse.ArgumentParser(description="Gemini API CLI with File & Context Caching") + parser.add_argument("-c", "--context", type=str, default=".gemini_context.json", + help="Path to the context file (defaults to .gemini_context.json)") + parser.add_argument("-f", "--files", nargs="+", default=[], + help="Files to upload to the Gemini API") + parser.add_argument("-m", "--model", type=str, default="models/gemini-1.5-pro-001", + help="The model to use (default: models/gemini-1.5-pro-001)") + parser.add_argument("-d", "--destroy", action="store_true", + help="Destroy the cloud files and cache, and delete the local context file") + parser.add_argument("-o", "--output", type=str, + help="Direct the raw output to a specific file instead of stdout") + parser.add_argument("-p", "--prompt", type=str, + help="The prompt to send to the AI") + parser.add_argument("positional_prompt", nargs=argparse.REMAINDER, + help="Positional arguments treated as the prompt if -p is omitted") + + args = parser.parse_args() + + # Determine the prompt text + prompt_text = args.prompt if args.prompt else " ".join(args.positional_prompt) if args.positional_prompt else None + + # Authenticate + api_key = os.environ.get("GEMINI_API_KEY") + if not api_key: + print("Error: GEMINI_API_KEY environment variable is not set.", file=sys.stderr) + sys.exit(1) + genai.configure(api_key=api_key) + + context_data = {"file_ids": [], "cache_id": None} + + # Load existing context if it exists + if os.path.exists(args.context): + try: + with open(args.context, "r") as f: + context_data = json.load(f) + except json.JSONDecodeError: + print(f"Warning: Could not parse {args.context}. Starting fresh.", file=sys.stderr) + + # --------------------------------------------------------- + # DESTROY FLAG LOGIC + # --------------------------------------------------------- + if args.destroy: + print("Destroying server resources and local context...") + if context_data.get("cache_id"): + try: + caching.CachedContent.get(context_data["cache_id"]).delete() + print(f"Deleted cache: {context_data['cache_id']}") + except Exception as e: + print(f"Warning: Failed to delete cache. {e}", file=sys.stderr) + + for file_id in context_data.get("file_ids", []): + try: + genai.delete_file(file_id) + print(f"Deleted file: {file_id}") + except Exception as e: + print(f"Warning: Failed to delete file '{file_id}'. {e}", file=sys.stderr) + + if os.path.exists(args.context): + os.remove(args.context) + print(f"Deleted local context file: {args.context}") + + print("Cleanup complete.") + return # Exit after destroying + + # --------------------------------------------------------- + # UPLOAD LOGIC + # --------------------------------------------------------- + if args.files: + for file_path in args.files: + if not os.path.exists(file_path): + print(f"Warning: File '{file_path}' not found. Skipping.", file=sys.stderr) + continue + + print(f"Uploading '{file_path}'...") + uploaded_file = genai.upload_file(file_path) + print(f"Success: '{file_path}' uploaded as '{uploaded_file.name}'") + + if uploaded_file.name not in context_data["file_ids"]: + context_data["file_ids"].append(uploaded_file.name) + + with open(args.context, "w") as f: + json.dump(context_data, f, indent=4) + + # --------------------------------------------------------- + # CACHE CREATION AND TTL EXTENSION LOGIC + # --------------------------------------------------------- + cache = None + system_instruction = ( + "You are a strict data extraction tool. Output exactly what is requested " + "(e.g., CSV). Never use markdown formatting blocks (like ```csv). " + "Never include conversational text, greetings, or explanations. Output raw data only." + ) + + if context_data.get("file_ids"): + if not context_data.get("cache_id"): + print("Creating Context Cache on Google's servers...") + file_objects = [genai.get_file(f_id) for f_id in context_data["file_ids"]] + + cache = caching.CachedContent.create( + model=args.model, + system_instruction=system_instruction, + contents=file_objects, + ttl=datetime.timedelta(minutes=60) + ) + context_data["cache_id"] = cache.name + with open(args.context, "w") as f: + json.dump(context_data, f, indent=4) + print(f"Context Cache created: {cache.name}") + else: + print(f"Loading existing cache: {context_data['cache_id']}") + cache = caching.CachedContent.get(context_data["cache_id"]) + print("Extending cache TTL by 60 minutes...") + cache.update(ttl=datetime.timedelta(minutes=60)) + + # --------------------------------------------------------- + # GENERATION LOGIC + # --------------------------------------------------------- + if prompt_text: + generation_config = { + "max_output_tokens": 65536, + "temperature": 0.0 + } + + if cache: + model = genai.GenerativeModel.from_cached_content( + cached_content=cache, + generation_config=generation_config + ) + else: + # Fallback if no files were uploaded + model = genai.GenerativeModel( + model_name=args.model, + system_instruction=system_instruction, + generation_config=generation_config + ) + + print("Generating response (this may take a moment for large outputs)...") + + response = model.generate_content(prompt_text) + + if args.output: + with open(args.output, "w") as f: + f.write(response.text) + print(f"Done! Raw output saved directly to {args.output}") + else: + print("-" * 40) + print(response.text) + print("-" * 40) + +if __name__ == "__main__": + main() \ No newline at end of file