diff options
Diffstat (limited to 'cli.py')
| -rwxr-xr-x | cli.py | 137 |
1 files changed, 91 insertions, 46 deletions
@@ -3,11 +3,12 @@ import argparse import os import sys +import json import ai_api_module #this is the CLI program used to interact with the GAI python module. -default_model=None +default_model = "gemini-2.0-flash-exp" def main(): """Main function to parse arguments and execute commands""" @@ -15,69 +16,95 @@ def main(): parser.add_argument( "--apikey", default=os.getenv("GEMINI_API_KEY"), help="API key for the service" ) - parser.add_argument("--modelid", default=default_model, help=f"ID of the model to use (default={default_model}") + parser.add_argument("--modelid", default=default_model, help=f"ID of the model to use (default={default_model})") parser.add_argument("--list-models", action="store_true", help="list the available models") - parser.add_argument("--assistantid", default="default_chatbot", help="ID of the assistant") - parser.add_argument("--list-assistants", "--list-a", "--list-assistant", action='store_true', help='list all the available assistants') - parser.add_argument('--list-conversations', '--list-conversation', '--list-convo', '--list-c', action='store_true', help='list previous conversations') + + # Conversation management + parser.add_argument("--new", "-n", action="store_true", help="Start a new conversation") + parser.add_argument("--save-conversation", metavar="FILE", help="Save conversation to file") + parser.add_argument("--load-conversation", metavar="FILE", help="Load conversation from file") + parser.add_argument("--list-conversations", action="store_true", help="List saved conversations") + parser.add_argument("--clear", action="store_true", help="Clear conversation history") + parser.add_argument("--system-instruction", metavar="TEXT", help="Set system instruction for the conversation") + + # Query parameters parser.add_argument("--query", "-q", default="", help="Query to be sent to the assistant") parser.add_argument("--topp", type=float, default=0.9, help="Top P value for the model") - parser.add_argument( - "--temperature", type=float, default=0.7, help="Temperature value for the model" - ) + parser.add_argument("--temperature", type=float, default=0.7, help="Temperature value for the model") + parser.add_argument("--max-tokens", type=int, default=8192, help="Maximum output tokens") parser.add_argument("--stream", action="store_true", help="Enable streaming response") parser.add_argument("--no-stream", action="store_true", help="Disable streaming response") - parser.add_argument("--output_mode", "-o", nargs=1, action="store", default='text', help="just output the json") + parser.add_argument("--output_mode", "-o", default='text', help="Output mode (text or json)") + parser.add_argument("--count-tokens", action="store_true", help="Count tokens in query without sending") - #parser.add_argument("--file-health", action="store_true", help="") - #parser.add_argument("--file-list", action="store_true", help="List all uploaded files") - #parser.add_argument("--file-upload", help="Upload a file to the service") - #parser.add_argument("--file-status", help="Check status of a file by ID") - #parser.add_argument("--file-purpose", default="assistants", help="Purpose of the file") - #parser.add_argument("--attach-files", nargs="*", help="File IDs to attach to the query") - - parser.add_argument("--new", "-n", help="force new session and new chat", action="store_true", default=False) #TODO: solidify a good approach here - parser.add_argument("extra_query", nargs="*") args = parser.parse_args() - # Try to get API key if not provided + # Get API key if args.apikey: ai_api_module.api_key = args.apikey else: ai_api_module.get_api_key() if not ai_api_module.api_key: - print("Error: API key not provided") + print("Error: API key not provided. Set GEMINI_API_KEY environment variable.") sys.exit(1) - if args.list_models: - res = ai_api_module.list_models(args.output_mode) - if res[0]: - for m in res[1]: - print(str(m)) - #print(f" ModelID: {m['id']}\n Description: {m['description']}\n Tokens: {m['max_completion_tokens']}\n Modalities: {str(m['modalities'])}") - print() + # Handle system instruction + if args.system_instruction: + ai_api_module.set_system_instruction(args.system_instruction) + print(f"System instruction set") + + # Handle conversation management + if args.new: + ai_api_module.clear_conversation() + print("Started new conversation") + + if args.load_conversation: + if ai_api_module.load_conversation(args.load_conversation): + print(f"Loaded conversation from {args.load_conversation}") else: - print('error: ') - print(str(res[1].text)) - return - + print(f"Failed to load conversation from {args.load_conversation}") + sys.exit(1) + if args.list_conversations: - res = ai_api_module.list_conversations(args.output_mode) - print(res) - return - - if args.list_assistants: - res = ai_api_module.list_assistants(args.output_mode) - if res[0]: - for a in res[1]: - print(f" ID: {a['id']}\n Name: {a['display_name']}") - print() + success, result = ai_api_module.list_saved_conversations() + if success: + if not result: + print("No saved conversations found") + else: + print("Saved conversations:") + for conv in result: + print(f" {conv['filename']}") + print(f" Saved: {conv['saved_at']}") + print(f" Messages: {conv['message_count']}") + print() else: - print('error: ' + str(res[1])) + print(f"Error: {result}") + return + + if args.clear: + ai_api_module.clear_conversation() + print("Conversation history cleared") + return + + # List models + if args.list_models: + success, result = ai_api_module.list_models(args.output_mode) + if success: + for model in result: + if args.output_mode == 'json': + print(json.dumps(model, indent=2)) + else: + print(f"Model: {model['display_name']}") + print(f" Name: {model['name']}") + print(f" Description: {model.get('description', 'N/A')}") + print(f" Input tokens: {model.get('input_token_limit', 'N/A')}") + print(f" Output tokens: {model.get('output_token_limit', 'N/A')}") + print() + else: + print(f'Error: {result}') return - # Process query if args.query or args.extra_query: @@ -87,17 +114,35 @@ def main(): query += " " + " ".join(args.extra_query) else: query = " ".join(args.extra_query) - + + # Count tokens if requested + if args.count_tokens: + success, count = ai_api_module.count_tokens(args.modelid, query) + if success: + print(f"Token count: {count}") + else: + print(f"Error counting tokens: {count}") + return + + # Determine streaming + use_stream = args.stream and not args.no_stream + response = ai_api_module.query( args.modelid, query, args.topp, args.temperature, - not args.no_stream and args.stream, + use_stream, + args.max_tokens, ) - if response and not args.stream: + if response and not use_stream: print(response) + + # Auto-save conversation if requested + if args.save_conversation: + filepath = ai_api_module.save_conversation(args.save_conversation) + print(f"\nConversation saved to {filepath}") else: parser.print_help() |