Tools: Qdrant Manual Upload Tool (Dynamic & Auto Collection)

gistfile1.txt

"""
title: Qdrant Manual Upload Tool (Dynamic & Auto Collection)
author: custom
description: Upload manual text or file content directly into Qdrant vector DB with dynamic collections.
version: 2.0.0
"""

import uuid
from typing import Optional, List

from pydantic import BaseModel, Field
from fastapi.concurrency import run_in_threadpool

from qdrant_client import QdrantClient
from qdrant_client.models import PointStruct, VectorParams, Distance

from open_webui.models.users import Users

from sentence_transformers import SentenceTransformer


# === CONFIG ===
QDRANT_URL = "http://localhost:6333"

# default fallback
DEFAULT_COLLECTION = "manual-default"


# === INIT ===
client = QdrantClient(url=QDRANT_URL)
model = SentenceTransformer("all-MiniLM-L6-v2")


# === HELPERS ===
async def _resolve_user(__user__):
    if not __user__ or not __user__.get("id"):
        raise ValueError("User required")
    return await run_in_threadpool(Users.get_user_by_id, str(__user__["id"]))


def _embed(text: str) -> List[float]:
    return model.encode(text).tolist()


def _ensure_collection(collection_name: str, vector_size: int):
    """
    Auto create collection if not exists
    """
    collections = client.get_collections().collections
    existing = [c.name for c in collections]

    if collection_name not in existing:
        client.create_collection(
            collection_name=collection_name,
            vectors_config=VectorParams(
                size=vector_size,
                distance=Distance.COSINE
            )
        )


def _chunk_text(text: str, chunk_size: int = 500, overlap: int = 50):
    """
    Simple chunking biar nggak 1 block gede
    """
    chunks = []
    start = 0

    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        start += chunk_size - overlap

    return chunks


# === TOOL CLASS ===
class Tools:

    class Valves(BaseModel):
        collection_name: str = Field(
            default=DEFAULT_COLLECTION,
            description="Default Qdrant collection"
        )
        chunk_size: int = 500
        chunk_overlap: int = 50

    def __init__(self):
        self.valves = self.Valves()

    async def upload_text(
        self,
        content: str,
        collection_name: Optional[str] = None,
        metadata: Optional[dict] = None,
        __user__: Optional[dict] = None,
    ) -> str:
        """
        Upload text (auto chunk + dynamic collection)

        Args:
            content: text
            collection_name: optional (override)
            metadata: optional
        """

        if not content.strip():
            return "Content kosong."

        try:
            user = await _resolve_user(__user__)
        except Exception as e:
            return f"User error: {e}"

        try:
            collection = collection_name or self.valves.collection_name

            # chunk text
            chunks = _chunk_text(
                content,
                self.valves.chunk_size,
                self.valves.chunk_overlap
            )

            # embed first chunk untuk tahu dimensi
            test_vector = _embed(chunks[0])

            # ensure collection
            _ensure_collection(collection, len(test_vector))

            points = []

            for chunk in chunks:
                vector = _embed(chunk)

                point_id = str(uuid.uuid4())

                payload = {
                    "content": chunk,
                    "user_id": str(user.id),
                    "collection": collection,
                    **(metadata or {})
                }

                points.append(
                    PointStruct(
                        id=point_id,
                        vector=vector,
                        payload=payload
                    )
                )

            client.upsert(
                collection_name=collection,
                points=points
            )

            return f"✅ Uploaded {len(points)} chunks ke collection '{collection}'"

        except Exception as e:
            return f"Upload gagal: {e}"

    async def search_text(
        self,
        query: str,
        collection_name: Optional[str] = None,
        limit: int = 5,
    ) -> str:
        """
        Search dari collection tertentu
        """

        if not query.strip():
            return "Query kosong."

        try:
            collection = collection_name or self.valves.collection_name

            vector = _embed(query)

            results = client.search(
                collection_name=collection,
                query_vector=vector,
                limit=limit
            )

            if not results:
                return "Tidak ada hasil."

            output = []
            for i, r in enumerate(results, 1):
                content = r.payload.get("content", "")
                score = r.score
                output.append(f"{i}. ({score:.4f}) {content[:200]}")

            return "\n".join(output)

        except Exception as e:
            return f"Search error: {e}"