Agentic LLM

1
import os
2
import uuid
3
import json
4
import requests
5
import sys
6
import time
7
8
API_KEY = os.getenv("ASI_ONE_API_KEY") or "sk-REPLACE_ME"
9
ENDPOINT = "https://api.asi1.ai/v1/chat/completions"
10
MODEL = "asi1"
11
TIMEOUT = 90 # seconds
12
13
# In-memory session management
14
15
SESSION_MAP: dict[str, str] = {}
16
17
def get_session_id(conv_id: str) -> str:
18
    """Return existing session UUID for this conversation or create a new one."""
19
    sid = SESSION_MAP.get(conv_id)
20
    if sid is None:
21
        sid = str(uuid.uuid4())
22
        SESSION_MAP[conv_id] = sid
23
    return sid
24
25
def ask(conv_id: str, messages: list[dict], *, stream: bool = False) -> str:
26
    """Send the messages list to the ASI:One agent and return the assistant reply."""
27
    session_id = get_session_id(conv_id)
28
    print(f"[session] Using session-id: {session_id}")
29
30
    headers = {
31
        "Authorization": f"Bearer {API_KEY}",
32
        "x-session-id": session_id,
33
        "Content-Type": "application/json",
34
    }
35
36
    payload = {
37
        "model": MODEL,
38
        "messages": messages,
39
        "stream": stream,
40
    }
41
42
    if not stream:
43
        resp = requests.post(ENDPOINT, headers=headers, json=payload, timeout=TIMEOUT)
44
        resp.raise_for_status()
45
        return resp.json()["choices"][0]["message"]["content"]
46
47
    # Streaming implementation
48
    with requests.post(ENDPOINT, headers=headers, json=payload, timeout=TIMEOUT, stream=True) as resp:
49
        resp.raise_for_status()
50
        full_text = ""
51
        for line in resp.iter_lines(decode_unicode=True):
52
            if not line or not line.startswith("data: "):
53
                continue
54
            line = line[len("data: ") :]
55
            if line == "[DONE]":
56
                break
57
            try:
58
                chunk = json.loads(line)
59
                choices = chunk.get("choices")
60
                if choices and "content" in choices[0].get("delta", {}):
61
                    token = choices[0]["delta"]["content"]
62
                    sys.stdout.write(token)
63
                    sys.stdout.flush()
64
                    full_text += token
65
            except json.JSONDecodeError:
66
                continue
67
        print()
68
        return full_text
69
70
if __name__ == "__main__": # Simple usage example
71
    conv_id = str(uuid.uuid4())
72
    messages = [
73
        {"role": "user", "content": "use Hi-dream model to generate image of monkey sitting on top of mountain"}
74
    ]
75
    reply = ask(conv_id, messages, stream=True)
76
    print(f"\nAssistant: {reply}")