Method search

src/langbot/pkg/vector/vdbs/pgvector_db.py:187–254

Search for similar vectors using cosine distance Args: collection: Collection name query_embedding: Query vector k: Number of top results to return Returns: Dictionary with search results in Chroma-compatible format

(
        self,
        collection: str,
        query_embedding: list[float],
        k: int = 5,
        search_type: str = 'vector',
        query_text: str = '',
        filter: dict[str, Any] | None = None,
        vector_weight: float | None = None,
    )

                raise

    async def search(
        self,
        collection: str,
        query_embedding: list[float],
        k: int = 5,
        search_type: str = 'vector',
        query_text: str = '',
        filter: dict[str, Any] | None = None,
        vector_weight: float | None = None,
    ) -> Dict[str, Any]:
        """Search for similar vectors using cosine distance

        Args:
            collection: Collection name
            query_embedding: Query vector
            k: Number of top results to return

        Returns:
            Dictionary with search results in Chroma-compatible format
        """
        await self.get_or_create_collection(collection)

        async with self.AsyncSessionLocal() as session:
            try:
                # Use cosine distance for similarity search
                from sqlalchemy import select

                # Query for similar vectors
                stmt = (
                    select(
                        PgVectorEntry.id,
                        PgVectorEntry.text,
                        PgVectorEntry.file_id,
                        PgVectorEntry.chunk_uuid,
                        PgVectorEntry.embedding.cosine_distance(query_embedding).label('distance'),
                    )
                    .filter(PgVectorEntry.collection == collection)
                    .order_by(PgVectorEntry.embedding.cosine_distance(query_embedding))
                    .limit(k)
                )

                if filter:
                    for cond in _build_pg_conditions(filter):
                        stmt = stmt.filter(cond)

                result = await session.execute(stmt)
                rows = result.fetchall()

                # Convert to Chroma-compatible format
                ids = []
                distances = []
                metadatas = []

                for row in rows:
                    ids.append(row.id)
                    distances.append(float(row.distance))
                    metadatas.append(
                        {'text': row.text or '', 'file_id': row.file_id or '', 'uuid': row.chunk_uuid or ''}