Haystack: интеграция с SpeShu.AI

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

from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.utils import Secret

generator = OpenAIChatGenerator(
    api_key=Secret.from_env_var("SPESHU_API_KEY", strict=True),
    api_base_url="https://speshu.ai/api/v1",
    model=os.environ["SPESHU_MODEL"],
    timeout=300,
    max_retries=2,
)

import os

from haystack import Document, Pipeline
from haystack.components.builders import ChatPromptBuilder
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
from haystack.dataclasses import ChatMessage
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.utils import Secret

BASE = "https://speshu.ai/api/v1"

document_store = InMemoryDocumentStore()
document_store.write_documents([
    Document(content="SpeShu.AI отдаёт OpenAI-совместимый API по адресу speshu.ai/api/v1."),
    Document(content="Стоимость запроса возвращается в usage.cost.total_cost в рублях."),
    Document(content="Потоковая передача включается параметром stream и завершается маркером [DONE]."),
])

template = [
    ChatMessage.from_system("Отвечай кратко и только по найденным документам."),
    ChatMessage.from_user(
        "Документы:\n"
        "{% for doc in documents %}{{ doc.content }}\n{% endfor %}\n"
        "Вопрос: {{ question }}"
    ),
]

pipeline = Pipeline()
pipeline.add_component("retriever", InMemoryBM25Retriever(document_store=document_store))
pipeline.add_component("prompt", ChatPromptBuilder(template=template))
pipeline.add_component(
    "llm",
    OpenAIChatGenerator(
        api_key=Secret.from_env_var("SPESHU_API_KEY", strict=True),
        api_base_url=BASE,
        model=os.environ["SPESHU_MODEL"],
    ),
)

pipeline.connect("retriever.documents", "prompt.documents")
pipeline.connect("prompt.prompt", "llm.messages")

question = "Где смотреть стоимость запроса?"
result = pipeline.run(
    {
        "retriever": {"query": question},
        "prompt": {"question": question},
    },
    {"llm": {"generation_kwargs": {"max_completion_tokens": 400}}},
)

reply = result["llm"]["replies"][0]
print(reply.text)
print("списано:", reply.meta["usage"]["cost"]["total_cost"], "₽")

import os

from haystack.components.embedders import OpenAITextEmbedder
from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.utils import Secret

# Эмбеддинги — отдельный провайдер, не SpeShu.AI.
EMBEDDINGS_API_BASE_URL = os.environ["OTHER_PROVIDER_BASE_URL"]
EMBEDDINGS_MODEL = os.environ["OTHER_PROVIDER_EMBEDDING_MODEL"]

document_store = InMemoryDocumentStore(embedding_dimensions=1536)
document_store.write_documents(docs)

embedder = OpenAITextEmbedder(
    api_key=Secret.from_env_var("OTHER_PROVIDER_API_KEY", strict=True),
    api_base_url=EMBEDDINGS_API_BASE_URL,
    model=EMBEDDINGS_MODEL,
)
retriever = InMemoryEmbeddingRetriever(document_store=document_store, embedder=embedder)

from haystack.dataclasses import StreamingChunk

def print_chunk(chunk: StreamingChunk) -> None:
    if chunk.content:
        print(chunk.content, end="", flush=True)

pipeline.run(data, include_outputs_from={"llm"}, streaming_callback=print_chunk)

from haystack.tools import Tool

weather = Tool(
    name="get_weather",
    description="Текущая погода в городе. Возвращает температуру в градусах Цельсия.",
    parameters={
        "type": "object",
        "properties": {"city": {"type": "string", "description": "Название города"}},
        "required": ["city"],
    },
    function=get_weather,
)

generator = OpenAIChatGenerator(..., tools=[weather])

import os
import time

import requests

BASE = "https://speshu.ai/api/v1"
HEADERS = {
    "Authorization": f"Bearer {os.environ['SPESHU_API_KEY']}",
    "Content-Type": "application/json",
}

created = requests.post(
    f"{BASE}/async/media/tasks",
    headers=HEADERS,
    json={"model": "nano-banana-2", "input": {"prompt": "Кот в смокинге, студийный свет"}},
    timeout=60,
)
created.raise_for_status()
task_id = created.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(task["resultJson"])
        break
    if task["state"] == "fail":
        raise RuntimeError(f"{task['failCode']}: {task['failMsg']}")
    time.sleep(3)