Try to get structured output from AI
This commit is contained in:
+82
-60
@@ -1,102 +1,124 @@
|
||||
import os
|
||||
import base64
|
||||
import json
|
||||
import urllib.request
|
||||
import urllib.error
|
||||
from typing import Any, Union
|
||||
from openai import OpenAI
|
||||
from pydantic import BaseModel
|
||||
|
||||
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"
|
||||
# Default base_url pointing to Lemonade's OpenAI-compatible endpoint
|
||||
API_HOST = os.getenv("OPENAI_API_HOST", "https://lemonade.matthiasg.dev/api/v1")
|
||||
AI_MODEL = os.getenv("OPENAI_MODEL", "LFM2.5-VL-1.6B-GGUF-Q8_0")
|
||||
AI_MODEL_THINK = os.getenv("OPENAI_MODEL_THINK", "False").lower() in ("true", "1", "t")
|
||||
OPENAPI_TOKEN = os.getenv("OPENAI_API_TOKEN", "lemonade")
|
||||
|
||||
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}")
|
||||
# Initialize OpenAI Client for Lemonade
|
||||
client = OpenAI(
|
||||
base_url=API_HOST,
|
||||
api_key=OPENAPI_TOKEN if OPENAPI_TOKEN else "lemonade"
|
||||
)
|
||||
|
||||
|
||||
def generate(
|
||||
model: str = AI_MODEL,
|
||||
prompt: str = "",
|
||||
images: list[str] = None,
|
||||
output_format: str = None,
|
||||
output_format: Union[str, dict, type[BaseModel]] = 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.
|
||||
Generate a completion response using the OpenAI SDK against Lemonade / OpenAI compatible APIs.
|
||||
Handles multimodal image content, structured outputs, and custom Lemonade parameters.
|
||||
"""
|
||||
if images is None:
|
||||
images = []
|
||||
|
||||
# Transform image file paths to base64 strings
|
||||
encoded_images = []
|
||||
messages = []
|
||||
|
||||
# 1. Add System Message if present
|
||||
if system_message:
|
||||
messages.append({"role": "system", "content": system_message})
|
||||
|
||||
# 2. Build User Content Payload (Text + Multimodal Images)
|
||||
user_content = []
|
||||
if prompt:
|
||||
user_content.append({"type": "text", "text": prompt})
|
||||
|
||||
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"))
|
||||
b64_str = 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)
|
||||
b64_str = img_path
|
||||
|
||||
body = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"images": encoded_images,
|
||||
# Ensure base64 string includes Data URI prefix for OpenAI vision format
|
||||
if not b64_str.startswith("data:"):
|
||||
image_url = f"data:image/jpeg;base64,{b64_str}"
|
||||
else:
|
||||
image_url = b64_str
|
||||
|
||||
user_content.append({
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_url}
|
||||
})
|
||||
|
||||
# Simplify content payload if text-only
|
||||
if len(user_content) == 1 and user_content[0]["type"] == "text":
|
||||
messages.append({"role": "user", "content": prompt})
|
||||
else:
|
||||
messages.append({"role": "user", "content": user_content})
|
||||
|
||||
# 3. Prepare parameters and custom body flags
|
||||
extra_body = {
|
||||
"think": should_think,
|
||||
"stream": False,
|
||||
}
|
||||
if not keep_alive:
|
||||
extra_body["keep_alive"] = "0m"
|
||||
|
||||
request_kwargs: dict[str, Any] = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"extra_body": extra_body,
|
||||
}
|
||||
|
||||
if system_message is not None:
|
||||
body["system"] = system_message
|
||||
|
||||
# 4. Handle Structured Outputs / JSON mode
|
||||
if output_format is not None:
|
||||
try:
|
||||
body["format"] = json.loads(output_format)
|
||||
except json.JSONDecodeError:
|
||||
body["format"] = output_format
|
||||
if isinstance(output_format, type) and issubclass(output_format, BaseModel):
|
||||
# Pydantic Model -> Use structured output parsing API
|
||||
response = client.beta.chat.completions.parse(
|
||||
response_format=output_format,
|
||||
**request_kwargs
|
||||
)
|
||||
return response.choices[0].message.content or ""
|
||||
|
||||
elif isinstance(output_format, dict):
|
||||
# JSON Schema dictionary
|
||||
request_kwargs["response_format"] = {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "structured_response",
|
||||
"strict": True,
|
||||
"schema": output_format,
|
||||
}
|
||||
}
|
||||
elif isinstance(output_format, str) and output_format.lower() == "json":
|
||||
# Standard JSON Mode
|
||||
request_kwargs["response_format"] = {"type": "json_object"}
|
||||
|
||||
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", "")
|
||||
# Standard completion request
|
||||
response = client.chat.completions.create(**request_kwargs)
|
||||
return response.choices[0].message.content or ""
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Example usage:
|
||||
result = generate(
|
||||
prompt="Explain what you see in the image or answer this prompt.",
|
||||
prompt="Explain what you see or answer this prompt.",
|
||||
should_think=AI_MODEL_THINK,
|
||||
)
|
||||
print(result)
|
||||
Reference in New Issue
Block a user