File size: 3,022 Bytes
f43f2d3 a811652 f43f2d3 a811652 f43f2d3 a811652 f43f2d3 7cbc718 f43f2d3 a811652 f43f2d3 a811652 f43f2d3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
"""
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
|