LlamaIndex

export SPESHU_API_KEY="sk-ваш-ключ"

export SPESHU_MODEL=$(curl -s https://speshu.ai/api/v1/models \
  -H "Authorization: Bearer $SPESHU_API_KEY" | jq -r '.data[0].id')
echo "$SPESHU_MODEL"

import os

import requests
from llama_index.core import Settings
from llama_index.llms.openai_like import OpenAILike

BASE = "https://speshu.ai/api/v1"
api_key = os.environ["SPESHU_API_KEY"]

catalog = requests.get(
    f"{BASE}/models",
    headers={"Authorization": f"Bearer {api_key}"},
    timeout=30,
).json()["data"]

entry = next(m for m in catalog if m["id"] == os.environ["SPESHU_MODEL"])

# context_window по умолчанию у OpenAILike — 3900 токенов, и из каталога он
# не подтягивается. Без явного значения контекст обрежется молча.
llm = OpenAILike(
    model=entry["id"],
    api_base=BASE,
    api_key=api_key,
    context_window=entry["context_length"],
    is_chat_model=True,
    is_function_calling_model=True,
)

Settings.llm = llm

import os

from llama_index.core import Document, Settings, VectorStoreIndex
from llama_index.embeddings.openai import OpenAIEmbedding

# --- отдельный провайдер эмбеддингов, не SpeShu.AI ---
embeddings = OpenAIEmbedding(
    model=os.environ["EMBEDDING_MODEL"],
    api_base=os.environ["EMBEDDINGS_API_BASE"],
    api_key=os.environ["EMBEDDINGS_API_KEY"],
)
Settings.embed_model = embeddings
# -------------------------------------

documents = [
    Document(text="SpeShu.AI даёт доступ к текстовым и мультимодальным моделям через один OpenAI-совместимый эндпоинт."),
    Document(text="Списание денег возвращается в поле usage.cost.total_cost в рублях."),
]

index = VectorStoreIndex.from_documents(documents)
query_engine = index.as_query_engine()

response = query_engine.query("Где смотрят списанную сумму?")
print(response)

from llama_index.embeddings.huggingface import HuggingFaceEmbedding

Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")

from llama_index.core.llms import ChatMessage

answer = Settings.llm.chat(
    [
        ChatMessage(role="system", content="Отвечай по-русски, кратко."),
        ChatMessage(role="user", content="Сравни потоковую и синхронную генерацию."),
    ]
)
print(answer.message.content)

streaming_engine = index.as_query_engine(streaming=True)

stream = streaming_engine.query("Напиши три абзаца про океан.")
for piece in stream.response_gen:
    print(piece, end="", flush=True)

import json

from llama_index.core.tools import FunctionTool


def get_weather(city: str) -> str:
    """Текущая погода в городе. Возвращает температуру в градусах Цельсия."""
    return json.dumps({"city": city, "temp_c": 14, "sky": "облачно"}, ensure_ascii=False)


weather_tool = FunctionTool.from_defaults(fn=get_weather)
tool_engine = index.as_query_engine(tools=[weather_tool])

print(tool_engine.query("Какая сейчас погода в Москве?"))

from llama_index.core.llms import ChatMessage
from pydantic import BaseModel


class Ticket(BaseModel):
    title: str
    priority: int
    tags: list[str]


structured_llm = Settings.llm.as_structured_llm(Ticket)
result = structured_llm.chat(
    [ChatMessage(role="user", content="Приоритет: высокий. Тема: не приходит письмо.")]
)
print(result.raw.model_dump_json())

import json
import os
import time

import requests

BASE = "https://speshu.ai/api/v1"
headers = {"Authorization": f"Bearer {os.environ['SPESHU_API_KEY']}"}

task_id = requests.post(
    f"{BASE}/async/media/tasks",
    headers={**headers, "Content-Type": "application/json"},
    json={"model": "nano-banana-2", "input": {"prompt": "Кот в смокинге, студийный свет"}},
    timeout=60,
).json()["data"]["taskId"]

while True:
    task = requests.get(
        f"{BASE}/async/media/tasks/{task_id}", headers=headers, timeout=30
    ).json()["data"]
    if task["state"] == "success":
        print(json.loads(task["resultJson"])["resultUrls"][0])
        break
    if task["state"] == "fail":
        raise RuntimeError(f"{task['failCode']}: {task['failMsg']}")
    time.sleep(3)