opensearch_setup.py

"""
Setup script for OpenSearch ML Commons models, pipelines, and connectors.
Converts curl commands to opensearch-py client calls.
"""

import os
import time
import json
from opensearchpy import OpenSearch
from opensearchpy.exceptions import RequestError


def wait_for_task(client, task_id, max_wait=300, poll_interval=2):
    """Wait for a task to complete."""
    print(f"Waiting for task {task_id} to complete...")
    start_time = time.time()
    
    while time.time() - start_time < max_wait:
        try:
            response = client.transport.perform_request(
                'GET',
                f'/_plugins/_ml/tasks/{task_id}'
            )
            
            state = response.get('state', '')
            if state == 'COMPLETED':
                print(f"Task {task_id} completed successfully")
                return response
            elif state == 'FAILED':
                error = response.get('error', 'Unknown error')
                raise Exception(f"Task {task_id} failed: {error}")
            
            time.sleep(poll_interval)
        except Exception as e:
            if 'not found' in str(e).lower():
                # Task might not be immediately available
                time.sleep(poll_interval)
                continue
            raise
    
    raise Exception(f"Task {task_id} did not complete within {max_wait} seconds")


def setup_opensearch():
    """Set up OpenSearch cluster with ML models, pipelines, and connectors."""
    
    # Initialize OpenSearch client
    client = OpenSearch(
        hosts=[{
            'host': os.getenv('OPENSEARCH_HOST', 'localhost'),
            'port': int(os.getenv('OPENSEARCH_PORT', 9200))
        }],
        use_ssl=False,
        verify_certs=False
    )
    
    # Store IDs and configuration
    config = {}
    
    print("=" * 60)
    print("Starting OpenSearch Setup")
    print("=" * 60)
    
    # Step 1: Configure Cluster Settings
    print("\n1. Configuring cluster settings...")
    try:
        cluster_settings = {
            "persistent": {
                "plugins.ml_commons.allow_registering_model_via_url": "true",
                "plugins.ml_commons.only_run_on_ml_node": "false",
                "plugins.ml_commons.model_access_control_enabled": "true",
                "plugins.ml_commons.native_memory_threshold": "99"
            }
        }
        response = client.cluster.put_settings(body=cluster_settings)
        print("✓ Cluster settings configured")
    except Exception as e:
        print(f"✗ Error configuring cluster settings: {e}")
        raise
    
    # Step 2: Register Model Group
    print("\n2. Registering model group...")
    try:
        model_group_body = {
            "name": "note_search_with_highlighter",
            "description": "Models for note search with highlighter"
        }
        response = client.transport.perform_request(
            'POST',
            '/_plugins/_ml/model_groups/_register',
            body=json.dumps(model_group_body),
            headers={'Content-Type': 'application/json'}
        )
        model_group_id = response.get('model_group_id')
        config['MODEL_GROUP_ID'] = model_group_id
        print(f"✓ Model group registered: {model_group_id}")
    except Exception as e:
        print(f"✗ Error registering model group: {e}")
        raise
    
    # Step 3: Register Embedding Model
    print("\n3. Registering embedding model...")
    try:
        embedding_model_body = {
            "name": "huggingface/sentence-transformers/all-MiniLM-L6-v2",
            "version": "1.0.1",
            "model_format": "TORCH_SCRIPT",
            "model_group_id": model_group_id
        }
        response = client.transport.perform_request(
            'POST',
            '/_plugins/_ml/models/_register',
            body=json.dumps(embedding_model_body),
            headers={'Content-Type': 'application/json'}
        )
        embedding_task_id = response.get('task_id')
        print(f"✓ Embedding model registration started: {embedding_task_id}")
        
        # Wait for model registration to complete
        task_response = wait_for_task(client, embedding_task_id)
        embedding_model_id = task_response.get('model_id')
        config['EMBEDDING_MODEL_ID'] = embedding_model_id
        print(f"✓ Embedding model registered: {embedding_model_id}")
    except Exception as e:
        print(f"✗ Error registering embedding model: {e}")
        raise
    
    # Step 4: Create Ingest Pipeline
    print("\n4. Creating ingest pipeline...")
    try:
        pipeline_body = {
            "description": "Generate embeddings for note content",
            "processors": [
                {
                    "text_embedding": {
                        "model_id": embedding_model_id,
                        "field_map": {
                            "content": "content_embedding"
                        }
                    }
                }
            ]
        }
        client.ingest.put_pipeline(id='note-embedding-pipeline', body=pipeline_body)
        config['INGEST_PIPELINE'] = 'note-embedding-pipeline'
        print("✓ Ingest pipeline created: note-embedding-pipeline")
    except Exception as e:
        print(f"✗ Error creating ingest pipeline: {e}")
        raise
    
    # Step 5: Create Notes Index
    print("\n5. Creating notes index...")
    try:
        index_body = {
            "settings": {
                "index": {
                    "number_of_shards": 1,
                    "number_of_replicas": 0,
                    "default_pipeline": "note-embedding-pipeline",
                    "knn": True
                },
                "analysis": {
                    "analyzer": {
                        "default": {
                            "type": "standard"
                        }
                    }
                }
            },
            "mappings": {
                "properties": {
                    "note_id": {
                        "type": "integer"
                    },
                    "title": {
                        "type": "text",
                        "analyzer": "standard"
                    },
                    "content": {
                        "type": "text",
                        "analyzer": "standard"
                    },
                    "text_to_embed": {
                        "type": "text"
                    },
                    "embedding": {
                        "type": "knn_vector",
                        "dimension": 384,
                        "method": {
                            "name": "hnsw",
                            "space_type": "cosinesimil",
                            "engine": "lucene",
                            "parameters": {
                                "ef_construction": 128,
                                "m": 24
                            }
                        }
                    },
                    "category": {
                        "type": "keyword"
                    },
                    "tags": {
                        "type": "keyword"
                    },
                    "user_id": {
                        "type": "integer"
                    },
                    "created_at": {
                        "type": "date"
                    },
                    "color": {
                        "type": "keyword"
                    }
                }
            }
        }
        
        # Delete index if it exists
        if client.indices.exists(index='notes'):
            client.indices.delete(index='notes')
            print("  Deleted existing notes index")
        
        client.indices.create(index='notes', body=index_body)
        config['INDEX_NAME'] = 'notes'
        print("✓ Notes index created")
    except RequestError as e:
        if 'resource_already_exists_exception' in str(e):
            print("  Index already exists, skipping...")
        else:
            print(f"✗ Error creating notes index: {e}")
            raise
    except Exception as e:
        print(f"✗ Error creating notes index: {e}")
        raise
    
    # Step 6: Register and Deploy Highlighting Model
    print("\n6. Registering highlighting model...")
    try:
        highlighting_model_body = {
            "name": "amazon/sentence-highlighting/opensearch-semantic-highlighter-v1",
            "version": "1.0.0",
            "model_format": "TORCH_SCRIPT",
            "function_name": "QUESTION_ANSWERING"
        }
        response = client.transport.perform_request(
            'POST',
            '/_plugins/_ml/models/_register?deploy=true',
            body=json.dumps(highlighting_model_body),
            headers={'Content-Type': 'application/json'}
        )
        highlighting_task_id = response.get('task_id')
        print(f"✓ Highlighting model registration started: {highlighting_task_id}")
        
        # Wait for model registration and deployment
        task_response = wait_for_task(client, highlighting_task_id)
        highlighting_model_id = task_response.get('model_id')
        config['HIGHLIGHTING_MODEL_ID'] = highlighting_model_id
        print(f"✓ Highlighting model registered and deployed: {highlighting_model_id}")
    except Exception as e:
        print(f"✗ Error registering highlighting model: {e}")
        raise
    
    # Step 7: Create Hybrid RRF Search Pipeline
    print("\n7. Creating hybrid RRF search pipeline...")
    try:
        rrf_pipeline_body = {
            "description": "Post processor for hybrid RRF search",
            "phase_results_processors": [
                {
                    "score-ranker-processor": {
                        "combination": {
                            "technique": "rrf"
                        }
                    }
                }
            ]
        }
        client.transport.perform_request(
            'PUT',
            '/_search/pipeline/hybrid-rrf-pipeline',
            body=json.dumps(rrf_pipeline_body),
            headers={'Content-Type': 'application/json'}
        )
        config['RRF_PIPELINE'] = 'hybrid-rrf-pipeline'
        print("✓ Hybrid RRF search pipeline created")
    except Exception as e:
        print(f"✗ Error creating RRF pipeline: {e}")
        raise
    
    # Step 8: Create DeepSeek Chat Connector
    print("\n8. Creating DeepSeek Chat connector...")
    try:
        connector_body = {
            "name": "DeepSeek Chat",
            "description": "Test connector for DeepSeek Chat",
            "version": "1",
            "protocol": "http",
            "parameters": {
                "endpoint": "api.deepseek.com",
                "model": "deepseek-chat"
            },
            "credential": {
                "deepSeek_key": os.getenv('DEEPSEEK_API_KEY', '')
            },
            "actions": [
                {
                    "action_type": "predict",
                    "method": "POST",
                    "url": "https://${parameters.endpoint}/v1/chat/completions",
                    "headers": {
                        "Content-Type": "application/json",
                        "Authorization": "Bearer ${credential.deepSeek_key}"
                    },
                    "request_body": "{ \"model\": \"${parameters.model}\", \"messages\": ${parameters.messages} }"
                }
            ]
        }
        response = client.transport.perform_request(
            'POST',
            '/_plugins/_ml/connectors/_create',
            body=json.dumps(connector_body),
            headers={'Content-Type': 'application/json'}
        )
        connector_id = response.get('connector_id')
        config['CONNECTOR_ID'] = connector_id
        print(f"✓ DeepSeek Chat connector created: {connector_id}")
    except Exception as e:
        print(f"✗ Error creating connector: {e}")
        raise
    
    # Step 9: Register Remote Model (DeepSeek Chat)
    print("\n9. Registering remote DeepSeek Chat model...")
    try:
        remote_model_body = {
            "name": "DeepSeek Chat model",
            "function_name": "remote",
            "description": "DeepSeek Chat",
            "model_group_id": model_group_id,
            "connector_id": connector_id
        }
        response = client.transport.perform_request(
            'POST',
            '/_plugins/_ml/models/_register?deploy=true',
            body=json.dumps(remote_model_body),
            headers={'Content-Type': 'application/json'}
        )
        remote_task_id = response.get('task_id')
        print(f"✓ Remote model registration started: {remote_task_id}")
        
        # Wait for remote model registration
        task_response = wait_for_task(client, remote_task_id)
        remote_model_id = task_response.get('model_id')
        config['REMOTE_MODEL_ID'] = remote_model_id
        print(f"✓ Remote model registered and deployed: {remote_model_id}")
    except Exception as e:
        print(f"✗ Error registering remote model: {e}")
        raise
    
    # Step 10: Create RAG Search Pipeline
    print("\n10. Creating RAG search pipeline...")
    try:
        rag_pipeline_body = {
            "response_processors": [
                {
                    "retrieval_augmented_generation": {
                        "tag": "Demo pipeline",
                        "description": "Demo pipeline Using DeepSeek Chat",
                        "model_id": remote_model_id,
                        "context_field_list": ["title", "content", "category", "tags"],
                        "system_prompt": "You are a helpful assistant that helps users explore and understand their personal notes. Use the provided context from the user's notes to answer questions accurately and helpfully. Reference specific notes by title when relevant. If the context doesn't contain enough information, acknowledge this politely. Be conversational and helpful.",
                        "user_instructions": "Answer based on these notes from my collection"
                    }
                }
            ]
        }
        client.transport.perform_request(
            'PUT',
            '/_search/pipeline/my-conversation-search-pipeline-deepseek-chat',
            body=json.dumps(rag_pipeline_body),
            headers={'Content-Type': 'application/json'}
        )
        config['CHATBOT_PIPELINE'] = 'my-conversation-search-pipeline-deepseek-chat'
        print("✓ RAG search pipeline created")
    except Exception as e:
        print(f"✗ Error creating RAG pipeline: {e}")
        raise
    
    print("\n" + "=" * 60)
    print("OpenSearch Setup Complete!")
    print("=" * 60)
    
    # Return environment variables
    env_vars = {
        'OPENSEARCH_HOST': os.getenv('OPENSEARCH_HOST', 'localhost'),
        'OPENSEARCH_PORT': os.getenv('OPENSEARCH_PORT', '9200'),
        **config
    }
    
    print("\nEnvironment Variables:")
    print("-" * 60)
    for key, value in env_vars.items():
        print(f"{key}={value}")
    print("-" * 60)
    
    return env_vars


if __name__ == '__main__':
    try:
        env_vars = setup_opensearch()
        print("\n✓ Setup completed successfully!")
        print("\nYou can now use these environment variables in your .env file:")
        for key, value in env_vars.items():
            print(f"{key}={value}")
    except Exception as e:
        print(f"\n✗ Setup failed: {e}")
        import traceback
        traceback.print_exc()
        exit(1)