Chroma

import os

import chromadb
from openai import OpenAI

# --- SpeShu.AI: генерация ---
speshu = OpenAI(
    base_url="https://speshu.ai/api/v1",
    api_key=os.environ["SPESHU_API_KEY"],
)

MODEL = os.environ["SPESHU_MODEL"]  # из GET /api/v1/models

# --- отдельный провайдер: только эмбеддинги, к SpeShu.AI отношения не имеет ---
embedder = OpenAI(
    base_url=os.environ["EMBEDDING_BASE_URL"],
    api_key=os.environ["EMBEDDING_API_KEY"],
)

EMBEDDING_MODEL = os.environ["EMBEDDING_MODEL"]


def embed(texts: list[str]) -> list[list[float]]:
    response = embedder.embeddings.create(model=EMBEDDING_MODEL, input=texts)
    return [item.embedding for item in response.data]


# --- локальное хранилище Chroma ---
client = chromadb.PersistentClient(path="./chroma")

import os

import chromadb
from openai import OpenAI

speshu = OpenAI(
    base_url="https://speshu.ai/api/v1",
    api_key=os.environ["SPESHU_API_KEY"],
)
MODEL = os.environ["SPESHU_MODEL"]

# отдельный провайдер эмбеддингов — не SpeShu.AI
embedder = OpenAI(
    base_url=os.environ["EMBEDDING_BASE_URL"],
    api_key=os.environ["EMBEDDING_API_KEY"],
)
EMBEDDING_MODEL = os.environ["EMBEDDING_MODEL"]

DOCUMENTS = [
    "Chroma хранит векторы, документы и метаданные в коллекциях.",
    "SpeShu.AI отдаёт генерацию текста по OpenAI-совместимому адресу.",
    "Списание денег видно в поле usage.cost.total_cost в рублях.",
]


def embed(texts: list[str]) -> list[list[float]]:
    return [item.embedding for item in embedder.embeddings.create(
        model=EMBEDDING_MODEL, input=texts
    ).data]


client = chromadb.PersistentClient(path="./chroma")
collection = client.get_or_create_collection("docs")

collection.add(
    ids=[f"doc-{index}" for index in range(len(DOCUMENTS))],
    documents=DOCUMENTS,
    embeddings=embed(DOCUMENTS),
    metadatas=[{"source": "handbook"} for _ in DOCUMENTS],
)

query = "Где смотреть потраченные деньги?"
hits = collection.query(query_embeddings=embed([query]), n_results=2)
context = hits["documents"][0]

answer = speshu.chat.completions.create(
    model=MODEL,
    messages=[
        {"role": "system", "content": "Отвечай только по приведённым фрагментам."},
        {
            "role": "user",
            "content": "Контекст:\n"
            + "\n".join(f"- {fragment}" for fragment in context)
            + f"\n\nВопрос: {query}",
        },
    ],
)

print(answer.choices[0].message.content)
print("списано:", answer.usage.cost.total_cost, "₽")

collection = client.get_or_create_collection("docs")
collection.add(
    ids=[f"doc-{index}" for index in range(len(DOCUMENTS))],
    documents=DOCUMENTS,
    metadatas=[{"source": "handbook", "lang": "ru"} for _ in DOCUMENTS],
)

matches = collection.get(where={"lang": "ru"})

from langchain_openai import ChatOpenAI

llm = ChatOpenAI(
    model=os.environ["SPESHU_MODEL"],
    base_url="https://speshu.ai/api/v1",
    api_key=os.environ["SPESHU_API_KEY"],
    temperature=0,
)