Spaces:
Sleeping
Sleeping
| """ | |
| Manages connections to external services: SSH, Solr, and Google Gemini. | |
| This module centralizes the initialization logic, making the main application | |
| cleaner and more focused on its primary task. It provides a single function | |
| to set up all necessary connections. | |
| """ | |
| import pysolr | |
| import google.generativeai as genai | |
| from sshtunnel import SSHTunnelForwarder | |
| import config | |
| def initialize_connections(): | |
| """ | |
| Establishes the SSH tunnel, and initializes Solr and Gemini clients. | |
| Returns: | |
| A tuple containing the initialized (ssh_tunnel_server, solr_client, llm_model). | |
| Returns (None, None, None) if any part of the initialization fails. | |
| """ | |
| ssh_tunnel_server = None | |
| try: | |
| # 1. Configure and start the SSH Tunnel | |
| print("Attempting to start SSH tunnel...") | |
| ssh_tunnel_server = SSHTunnelForwarder( | |
| (config.SSH_HOST, config.SSH_PORT), | |
| ssh_username=config.SSH_USER, | |
| ssh_password=config.SSH_PASS, | |
| remote_bind_address=(config.REMOTE_SOLR_HOST, config.REMOTE_SOLR_PORT), | |
| local_bind_address=('127.0.0.1', config.LOCAL_BIND_PORT) | |
| ) | |
| ssh_tunnel_server.start() | |
| print(f"🚀 SSH tunnel established: Local Port {ssh_tunnel_server.local_bind_port} -> Remote Solr.") | |
| # 2. Initialize the pysolr client | |
| solr_url = f'http://127.0.0.1:{ssh_tunnel_server.local_bind_port}/solr/{config.SOLR_CORE_NAME}' | |
| solr_client = pysolr.Solr(solr_url, auth=(config.SOLR_USER, config.SOLR_PASS), always_commit=True) | |
| solr_client.ping() # Verify connection | |
| print(f"✅ Solr connection successful on core '{config.SOLR_CORE_NAME}'.") | |
| # 3. Initialize the LLM | |
| if not config.GEMINI_API_KEY: | |
| print("❌ CRITICAL: GEMINI_API_KEY is not set. LLM will not be available.") | |
| raise ValueError("GEMINI_API_KEY is missing.") | |
| genai.configure(api_key=config.GEMINI_API_KEY) | |
| llm_model = genai.GenerativeModel('gemini-2.5-flash', generation_config=genai.types.GenerationConfig(temperature=0)) | |
| print(f"✅ LLM Model '{llm_model.model_name}' initialized.") | |
| print("✅ System Initialized Successfully.") | |
| return ssh_tunnel_server, solr_client, llm_model | |
| except pysolr.SolrError as e: | |
| print(f"\n❌ Solr Error: {e}") | |
| print("Please check if the Solr core '{config.SOLR_CORE_NAME}' exists and is running.") | |
| if ssh_tunnel_server and ssh_tunnel_server.is_active: | |
| ssh_tunnel_server.stop() | |
| return None, None, None | |
| except ValueError as e: | |
| print(f"\n❌ Configuration Error: {e}") | |
| if ssh_tunnel_server and ssh_tunnel_server.is_active: | |
| ssh_tunnel_server.stop() | |
| return None, None, None | |
| except Exception as e: | |
| print(f"\n❌ An unexpected error occurred during setup: {e}") | |
| if ssh_tunnel_server and ssh_tunnel_server.is_active: | |
| ssh_tunnel_server.stop() | |
| return None, None, None | |