ai-time-machine / src /time_machine /adapters /llm /cloud_completion.py
manikandanj's picture
Prepare AI Time Machine hackathon Space
5862322 verified
Raw
History Blame Contribute Delete
8.18 kB
"""Cloud completion function for OpenAI-compatible APIs.
Works with Together AI, OpenRouter, Fireworks, OpenAI, and any provider
that speaks the ``/v1/chat/completions`` protocol. Uses only stdlib so
it adds zero new dependencies.
"""
from __future__ import annotations
import json
import logging
import os
import urllib.error
import urllib.request
from collections.abc import Callable, Iterator
from typing import Any
from time_machine.domain.errors import AdapterConfigurationError
CompletionFn = Callable[[str, int], str]
StreamCompletionFn = Callable[[str, int], Iterator[str]]
logger = logging.getLogger(__name__)
def create_cloud_completion_fn(
api_key: str | None = None,
base_url: str | None = None,
model: str | None = None,
) -> CompletionFn:
"""Return a ``(prompt, max_new_tokens) -> str`` callable for the Qwen adapter.
Configuration priority: explicit args > env vars > defaults.
Env vars:
TIME_MACHINE_LLM_API_KEY - required, unless TOGETHER_API_KEY is set
TOGETHER_API_KEY - fallback key for Together AI
TIME_MACHINE_LLM_BASE_URL - default: ``https://api.together.xyz/v1``
TIME_MACHINE_LLM_MODEL - default: ``Qwen/Qwen2.5-7B-Instruct-Turbo``
"""
key = api_key or os.getenv("TIME_MACHINE_LLM_API_KEY") or os.getenv("TOGETHER_API_KEY", "")
url = (
base_url
or os.getenv("TIME_MACHINE_LLM_BASE_URL", "https://api.together.xyz/v1")
).rstrip("/")
mdl = model or os.getenv("TIME_MACHINE_LLM_MODEL", "Qwen/Qwen2.5-7B-Instruct-Turbo")
if not key:
raise AdapterConfigurationError(
"TIME_MACHINE_LLM_API_KEY is required for the cloud LLM profile. "
"Set it to your Together AI / OpenRouter / OpenAI API key, or set "
"TOGETHER_API_KEY for Together AI."
)
logger.info("Cloud LLM: model=%s base_url=%s", mdl, url)
def complete(prompt: str, max_new_tokens: int) -> str:
payload: dict[str, Any] = {
"model": mdl,
"messages": [
{
"role": "system",
"content": (
"You are a structured-output assistant. "
"Return ONLY the requested JSON object, no prose."
),
},
{"role": "user", "content": prompt},
],
"max_tokens": max_new_tokens,
"temperature": 0.75,
"response_format": {"type": "json_object"},
}
result = _post_chat_completion(url=url, api_key=key, payload=payload)
try:
content = result["choices"][0]["message"]["content"]
except (KeyError, IndexError, TypeError) as exc:
raise AdapterConfigurationError(
f"Cloud LLM API returned an unexpected payload: {result!r}"
) from exc
logger.debug("Cloud LLM raw response length: %d chars", len(content))
return content
return complete
def create_cloud_stream_completion_fn(
api_key: str | None = None,
base_url: str | None = None,
model: str | None = None,
) -> StreamCompletionFn:
"""Return a streaming ``(prompt, max_new_tokens) -> text chunks`` callable.
This uses the OpenAI-compatible chat completions SSE protocol. It is intended
for spoken conversation turns, not the structured JSON generation steps.
"""
key = api_key or os.getenv("TIME_MACHINE_LLM_API_KEY") or os.getenv("TOGETHER_API_KEY", "")
url = (
base_url
or os.getenv("TIME_MACHINE_LLM_BASE_URL", "https://api.together.xyz/v1")
).rstrip("/")
mdl = model or os.getenv("TIME_MACHINE_LLM_MODEL", "Qwen/Qwen2.5-7B-Instruct-Turbo")
if not key:
raise AdapterConfigurationError(
"TIME_MACHINE_LLM_API_KEY is required for the cloud LLM profile. "
"Set it to your Together AI / OpenRouter / OpenAI API key, or set "
"TOGETHER_API_KEY for Together AI."
)
def stream(prompt: str, max_new_tokens: int) -> Iterator[str]:
payload: dict[str, Any] = {
"model": mdl,
"messages": [
{
"role": "system",
"content": (
"You are a concise in-character voice actor. "
"Return only the spoken reply text, no JSON, no labels."
),
},
{"role": "user", "content": prompt},
],
"max_tokens": max_new_tokens,
"temperature": 0.75,
"stream": True,
}
yield from _post_chat_completion_stream(url=url, api_key=key, payload=payload)
return stream
def _post_chat_completion(
url: str,
api_key: str,
payload: dict[str, Any],
) -> dict[str, Any]:
body = json.dumps(payload).encode("utf-8")
req = urllib.request.Request(
f"{url}/chat/completions",
data=body,
headers={
"Authorization": f"Bearer {api_key}",
"Accept": "application/json",
"Content-Type": "application/json",
"User-Agent": "ai-time-machine/0.1 urllib",
},
method="POST",
)
try:
with urllib.request.urlopen(req, timeout=60) as resp:
return json.loads(resp.read().decode("utf-8"))
except urllib.error.HTTPError as exc:
error_body = exc.read().decode("utf-8", errors="replace")
raise AdapterConfigurationError(
f"Cloud LLM API returned {exc.code}: {error_body}"
) from exc
except urllib.error.URLError as exc:
raise AdapterConfigurationError(
f"Cloud LLM API request failed: {exc.reason}"
) from exc
except TimeoutError as exc:
raise AdapterConfigurationError("Cloud LLM API request timed out.") from exc
def _post_chat_completion_stream(
url: str,
api_key: str,
payload: dict[str, Any],
) -> Iterator[str]:
body = json.dumps(payload).encode("utf-8")
req = urllib.request.Request(
f"{url}/chat/completions",
data=body,
headers={
"Authorization": f"Bearer {api_key}",
"Accept": "text/event-stream",
"Content-Type": "application/json",
"User-Agent": "ai-time-machine/0.1 urllib",
},
method="POST",
)
try:
with urllib.request.urlopen(req, timeout=60) as resp:
for raw_line in resp:
line = raw_line.decode("utf-8", errors="replace").strip()
if not line or not line.startswith("data:"):
continue
data = line.removeprefix("data:").strip()
if data == "[DONE]":
return
chunk = _decode_stream_chunk(data)
if chunk:
yield chunk
except urllib.error.HTTPError as exc:
error_body = exc.read().decode("utf-8", errors="replace")
raise AdapterConfigurationError(
f"Cloud LLM streaming API returned {exc.code}: {error_body}"
) from exc
except urllib.error.URLError as exc:
raise AdapterConfigurationError(
f"Cloud LLM streaming API request failed: {exc.reason}"
) from exc
except TimeoutError as exc:
raise AdapterConfigurationError("Cloud LLM streaming API request timed out.") from exc
def _decode_stream_chunk(data: str) -> str:
try:
decoded = json.loads(data)
except json.JSONDecodeError:
return ""
try:
choice = decoded["choices"][0]
except (KeyError, IndexError, TypeError):
return ""
delta = choice.get("delta") if isinstance(choice, dict) else None
if isinstance(delta, dict) and isinstance(delta.get("content"), str):
return delta["content"]
text = choice.get("text") if isinstance(choice, dict) else None
return text if isinstance(text, str) else ""