diff options
Diffstat (limited to 'ai_api_module.py')
| -rwxr-xr-x | ai_api_module.py | 183 |
1 files changed, 183 insertions, 0 deletions
diff --git a/ai_api_module.py b/ai_api_module.py new file mode 100755 index 0000000..16f80df --- /dev/null +++ b/ai_api_module.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python +# COPYRIGHT 2025 Thomas Grothe +import argparse +import requests +import json +import os +import sys +import http.client +from util import * + +#this is a python module to provide REST API access to some AI model +# use a CLI tool or custom webtool to use it +# docs: TODO + +# each function shall return a tuple containing [success bool, response dict] + +### CONFIGURATION HERE (yes you could move this into a config file if desired, but i don't need that currently) +# Base URL and endpoints configuration + +# the key of the environment variable used to store your api key +api_key_env_key = 'GEMINI_API_KEY' + +# the api key, if you prefer to set it here +api_key = None + +urlbase = 'https://generativelanguage.googleapis.com/v1beta/' #models/gemini-2.0-flash:generateContent' +endpoints = { + 'chat': { + 'simple': urlbase+'/chat/completions', # without custom assistant #use components/schemas/ChatCompletionRequest + 'assistant': urlbase+'/assistant/--assistant_id--/chat/completions', #use components/schemas/AssistantChatCompletionRequest + }, + 'model': { + 'list': urlbase+'/models', + }, + 'assistant': { + 'list': urlbase+'/assistants' + }, + 'conversation': { + 'list': urlbase+'/conversations', + 'get': urlbase+'/conversation/' #/message_id + }, + 'genContent': 'models/gemini-2.0-flash:generateContent?key=' +} + +# Possible environment variables for API key + +#the default message context. this will be appended to or replaced as conversation goes, or if using a previous chat +messages = [ + { + 'role': 'system', + 'content': '', #TODO system prompt instructions + }, +] + +def setupURLParams(): + urlbase += f'?key={api_key}' + +def get_api_key(): + global api_key + """Try to get API key from environment variables""" + if api_key == None: + if os.environ.get(api_key_env_key): + api_key = os.environ.get(api_key_env_key) + return api_key + return None + +def list_models(output_mode = ''): + headers = {"Content-Type": "application/json"} + url = endpoints['model']['list'] + f'?key={api_key}' + print(f'calling {url}') + response = requests.get(f"{endpoints['model']['list']}", headers=headers) + models = [] + + if response.status_code == 200: + j = json.loads(response.text) + if 'j' in output_mode: + return j + else: + for d in j['data']: + models.append(d) + return [True, models] + else: + return [False, response] + +def list_assistants(output_mode = ''): + headers = {"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"} + response = requests.get(f"{endpoints['assistant']['list']}", headers=headers) + if response.status_code == 200: + j = json.loads(response.text) # keys: assistants, total_assistants, scores,carousel_display_name,carousel_subtitle + if 'j' in output_mode: + return j + else: + res = [] + for a in j['assistants']: + res.append(a) + return [True, res] + else: + return [False, response] + +def list_conversations(output_mode = ''): + headers = {"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"} + response = requests.get(f"{endpoints['conversation']['list']}", headers=headers) + if response.status_code == 200: + j = json.loads(response.text) + if 'j' in output_mode: + return j + else: + res = [] + for a in j.keys(): #iterate through assistants + for cid in a: #j[]: #convo ids for each assistant + print('todo') + #TODO this gives only the first message, with its id, then we have to call /conversation/{m_id} + + return [True, res] + +def get_conversations(msg_id, output_mode = ''): + headers = {"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"} + response = requests.get(f"{endpoints['conversation']['get']}", headers=headers) + + +def query(model_id, query, top_p, temperature, stream=False): + """Call the model with some query""" + headers = {"Content-Type": "application/json"} #, "Authorization": f"Bearer {api_key}"} + url = urlbase+endpoints['genContent']+str(api_key) + #attachments = [] #file attachments + + payload = { #todo options for other parameters + "contents": [ + { + "parts": [ + { + "text": query + } + ] + } + ] + } + # payload = { + # "model": model_id, + # "top_p": top_p, + # "temperature": temperature, + # "stream": stream, + # "messages": [{"role": "user", "content": query}], #TODO include previous conversation here + # # components/schemas/ChatCompletionsMessage-Input + # # content: string + # # attachments: + # # role: user, assistant, or system + # # name: + # # tool_call_id + # # tool_calls + + #} + + if stream: + response = requests.post(url, json=payload, headers=headers, stream=True) + for line in response.iter_lines(): + if line: + line = line.decode('utf-8') + if line.startswith("data: ") and not line.startswith("data: [DONE]"): + data = line[len("data: "):] + try: + content = json.loads(data)["choices"][0]["delta"]["content"] + if content: + log(content) + sys.stdout.write(content) + sys.stdout.flush() + except (json.JSONDecodeError, KeyError, IndexError): + pass + print() + else: + response = requests.post(url, json=payload, headers=headers) + if response.status_code == 200: + response_json = response.json() + #TODO handle message to append + response = response_json['candidates'][0]['content']['parts'][0]['text'] + log(response) + return response + #return response_json["choices"][0]["message"]["content"] + else: + print(f"Error: {response.status_code}") + print(response.text) + log(response.text) + return None |