# NIKAME GENERATED — DO NOT EDIT DIRECTLY
from qdrant_client import AsyncQdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct
{% if has_ollama %}
import httpx
from app.core.config import settings
{% else %}
from sentence_transformers import SentenceTransformer
{% endif %}

client = AsyncQdrantClient("http://qdrant:6333")
{% if not has_ollama %}
model = SentenceTransformer('all-MiniLM-L6-v2')
{% endif %}

def get_collection_name(table_name: str):
    return f"nikame_{table_name}"

async def get_embedding(text: str) -> list[float]:
    {% if has_ollama %}
    url = f"{settings.OLLAMA_HOST}/api/embeddings"
    try:
        async with httpx.AsyncClient(timeout=10.0) as http_client:
            response = await http_client.post(
                url,
                json={"model": "nomic-embed-text", "prompt": text}
            )
            response.raise_for_status()
            return response.json()["embedding"]
    except (httpx.ConnectError, httpx.ReadTimeout):
        from fastapi import HTTPException
        raise HTTPException(
            status_code=503,
            detail={"error": "model_loading", "retry_after": 30}
        )
    {% else %}
    return model.encode(text).tolist()
    {% endif %}

async def ensure_collection(collection_name: str, vector_size: int = {% if has_ollama %}768{% else %}384{% endif %}):
    collections = await client.get_collections()
    exists = any(c.name == collection_name for c in collections.collections)
    if not exists:
        await client.create_collection(
            collection_name=collection_name,
            vectors_config=VectorParams(size=vector_size, distance=Distance.COSINE),
        )

async def upsert_vector(collection_name: str, row_id: str, text: str, payload: dict):
    vector = await get_embedding(text)
    await client.upsert(
        collection_name=collection_name,
        points=[
            PointStruct(id=row_id, vector=vector, payload=payload)
        ]
    )

async def search_vector(collection_name: str, query_text: str, limit: int = 5):
    query_vector = await get_embedding(query_text)
    return await client.search(
        collection_name=collection_name,
        query_vector=query_vector,
        limit=limit
    )
