File size: 2,257 Bytes
f43f2d3 6f4b6a3 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 |
"""
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
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()
print(f"✅ Solr connection successful on core '{config.SOLR_CORE_NAME}'.")
# 3. Initialize the LLM
genai.configure(api_key=config.GEMINI_API_KEY)
llm_model = genai.GenerativeModel('gemini-2.5-flash', generation_config=genai.types.GenerationConfig(temperature=0, max_output_tokens=1048576))
print(f"✅ LLM Model '{llm_model.model_name}' initialized.")
print("✅ System Initialized Successfully.")
return ssh_tunnel_server, solr_client, llm_model
except Exception as e:
print(f"\n❌ An error occurred during setup: {e}")
if ssh_tunnel_server and ssh_tunnel_server.is_active:
ssh_tunnel_server.stop()
return None, None, None
|