102 lines
3.2 KiB
Python
102 lines
3.2 KiB
Python
import os
|
|
import base64
|
|
import json
|
|
import urllib.request
|
|
import urllib.error
|
|
|
|
API_HOST = os.getenv("OPENAI_API_HOST", "https://chat.matthiasg.dev/ollama")
|
|
AI_MODEL = os.getenv("OPENAI_MODEL", "llava:7b-v1.6-mistral-q4_1")
|
|
AI_MODEL_THINK = os.getenv("OPENAI_MODEL_THINK", "True").lower() in ("true", "1", "t")
|
|
OPENAPI_TOKEN = os.getenv("OPENAI_API_TOKEN", None)
|
|
OPENAPI_ENDPOINT = "/api/generate"
|
|
|
|
print(f"Using API Host: {API_HOST}")
|
|
print(f"Using API Model: {AI_MODEL}")
|
|
print(f"Using API Model Think: {AI_MODEL_THINK}")
|
|
print(f"Using API Token: {'Yes' if OPENAPI_TOKEN else 'No'} {OPENAPI_TOKEN[:5] + '...' if OPENAPI_TOKEN else ''}")
|
|
|
|
def call_api(body: dict, endpoint: str = OPENAPI_ENDPOINT) -> str:
|
|
"""Call the API with the given endpoint and body dict."""
|
|
url = f"{API_HOST}{endpoint}"
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
}
|
|
if OPENAPI_TOKEN:
|
|
headers["Authorization"] = f"Bearer {OPENAPI_TOKEN}"
|
|
|
|
json_data = json.dumps(body).encode("utf-8")
|
|
req = urllib.request.Request(url, data=json_data, headers=headers, method="POST")
|
|
|
|
try:
|
|
with urllib.request.urlopen(req) as response:
|
|
return response.read().decode("utf-8")
|
|
except urllib.error.HTTPError as e:
|
|
error_body = e.read().decode("utf-8")
|
|
raise Exception(f"Error calling API: HTTP {e.code} - {error_body}")
|
|
except urllib.error.URLError as e:
|
|
raise Exception(f"Failed to reach server: {e.reason}")
|
|
|
|
|
|
def generate(
|
|
model: str = AI_MODEL,
|
|
prompt: str = "",
|
|
images: list[str] = None,
|
|
output_format: str = None,
|
|
system_message: str = None,
|
|
keep_alive: bool = True,
|
|
should_think: bool = AI_MODEL_THINK,
|
|
) -> str:
|
|
"""
|
|
Generate a response for a given prompt with a provided model via the Ollama/OpenAI API.
|
|
Handles base64 encoding for local image file paths and structures the request body.
|
|
"""
|
|
if images is None:
|
|
images = []
|
|
|
|
# Transform image file paths to base64 strings
|
|
encoded_images = []
|
|
for img_path in images:
|
|
if os.path.isfile(img_path):
|
|
with open(img_path, "rb") as image_file:
|
|
encoded_images.append(base64.b64encode(image_file.read()).decode("utf-8"))
|
|
else:
|
|
# If it's already a base64 string or an invalid path, keep as-is
|
|
encoded_images.append(img_path)
|
|
|
|
body = {
|
|
"model": model,
|
|
"prompt": prompt,
|
|
"images": encoded_images,
|
|
"think": should_think,
|
|
"stream": False,
|
|
}
|
|
|
|
if system_message is not None:
|
|
body["system"] = system_message
|
|
|
|
if output_format is not None:
|
|
try:
|
|
body["format"] = json.loads(output_format)
|
|
except json.JSONDecodeError:
|
|
body["format"] = output_format
|
|
|
|
if not keep_alive:
|
|
body["keep_alive"] = "0m"
|
|
|
|
response_text = call_api(body)
|
|
|
|
try:
|
|
decoded_response = json.loads(response_text)
|
|
except json.JSONDecodeError as e:
|
|
raise Exception(f"Error decoding JSON response: {e}")
|
|
|
|
return decoded_response.get("response", "")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Example usage:
|
|
result = generate(
|
|
prompt="Explain what you see in the image or answer this prompt.",
|
|
should_think=AI_MODEL_THINK,
|
|
)
|
|
print(result) |