Class DgraphVectorClient

vector/beir/dgraph_client.py:12–277

Client for interacting with Dgraph for vector search.

class DgraphVectorClient:
    """Client for interacting with Dgraph for vector search."""
    
    def __init__(self, config: DgraphConfig):
        """
        Initialize Dgraph client.
        
        Args:
            config: Dgraph configuration
        """
        self.config = config
        
        # Increase gRPC buffer limits for large vector batches
        # Default is 4MB, increase to 256MB for batch inserts
        grpc_options = [
            ('grpc.max_send_message_length', 256 * 1024 * 1024),
            ('grpc.max_receive_message_length', 256 * 1024 * 1024),
        ]
        self.client_stub = pydgraph.DgraphClientStub(
            addr=config.grpc_endpoint,
            options=grpc_options
        )
        self.client = pydgraph.DgraphClient(self.client_stub)
        print(f"Connected to Dgraph at {config.grpc_endpoint} (256MB gRPC buffers)")
    
    def setup_schema(self, hnsw_config: HNSWConfig):
        """
        Set up the Dgraph schema for BEIR documents.
        
        Args:
            hnsw_config: HNSW index configuration
        """
        print("Setting up Dgraph schema...")
        
        schema = f"""
            doc_id: string @index(hash) .
            title: string @index(term, fulltext) .
            text: string @index(fulltext) .
            embedding: float32vector @index({hnsw_config.to_index_string()}) .
            
            type Document {{
                doc_id
                title
                text
                embedding
            }}
        """
        
        op = pydgraph.Operation(schema=schema)
        self.client.alter(op)
        print("Schema setup complete")
    
    def reset_database(self):
        """Drop all data and schema."""
        print("Dropping all data and schema...")
        op = pydgraph.Operation(drop_all=True)
        self.client.alter(op)
        print("Database reset complete")