105 lines
3.5 KiB
Python
105 lines
3.5 KiB
Python
"""
|
|
Database module for the application.
|
|
Handles connections to PostgreSQL (Zulip DB) and ChromaDB.
|
|
"""
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import sessionmaker, scoped_session
|
|
from sqlalchemy.ext.declarative import declarative_base
|
|
import chromadb
|
|
|
|
# SQLAlchemy base class for models
|
|
Base = declarative_base()
|
|
|
|
# Global variables for SQLAlchemy
|
|
db_engine = None
|
|
db_session = None
|
|
|
|
# Global variable for ChromaDB
|
|
chroma_client = None
|
|
chroma_collection = None
|
|
|
|
def init_db(app):
|
|
"""
|
|
Initialize database connections.
|
|
|
|
Args:
|
|
app: Flask application object
|
|
"""
|
|
global db_engine, db_session, chroma_client, chroma_collection
|
|
|
|
# Initialize SQLAlchemy engine and session
|
|
db_engine = create_engine(app.config['SQLALCHEMY_DATABASE_URI'])
|
|
db_session = scoped_session(sessionmaker(autocommit=False, autoflush=False, bind=db_engine))
|
|
|
|
# Set query property for models
|
|
Base.query = db_session.query_property()
|
|
|
|
# Initialize ChromaDB
|
|
try:
|
|
# Set allow_reset to True to prevent "Add of existing embedding ID" warnings
|
|
chroma_client = chromadb.PersistentClient(
|
|
path=app.config['CHROMADB_PATH'],
|
|
settings=chromadb.Settings(
|
|
allow_reset=True,
|
|
anonymized_telemetry=False,
|
|
is_persistent=True
|
|
)
|
|
)
|
|
|
|
# Import here to avoid circular imports
|
|
from app.db.chroma_service import CustomEmbeddingFunction
|
|
|
|
# Create embedding function with setting from config
|
|
try:
|
|
# Always use Ollama since it's more reliable
|
|
embedding_function = CustomEmbeddingFunction(use_nomic=False)
|
|
|
|
# Get or create ChromaDB collection for Zulip messages with custom embedding function
|
|
chroma_collection = chroma_client.get_or_create_collection(
|
|
name=app.config.get('CHROMADB_COLLECTION', 'zulip_messages'),
|
|
metadata={
|
|
"hnsw:space": "cosine",
|
|
"hnsw:allow_replace_deleted": True # Allow replacing deleted vectors
|
|
},
|
|
embedding_function=embedding_function
|
|
)
|
|
except Exception as e:
|
|
print(f"Error with embedding function: {e}")
|
|
print("Creating collection without embedding function")
|
|
# Create collection without embedding function
|
|
chroma_collection = chroma_client.get_or_create_collection(
|
|
name=app.config.get('CHROMADB_COLLECTION', 'zulip_messages'),
|
|
metadata={
|
|
"hnsw:space": "cosine",
|
|
"hnsw:allow_replace_deleted": True # Allow replacing deleted vectors
|
|
}
|
|
)
|
|
except Exception as e:
|
|
print(f"Critical error initializing ChromaDB: {e}")
|
|
print("ChromaDB functionality will not be available")
|
|
chroma_client = None
|
|
chroma_collection = None
|
|
|
|
# Register teardown function to remove database sessions
|
|
@app.teardown_appcontext
|
|
def shutdown_session(exception=None):
|
|
"""Remove the database session at the end of the request."""
|
|
db_session.remove()
|
|
|
|
def get_db_session():
|
|
"""
|
|
Get the current database session.
|
|
|
|
Returns:
|
|
SQLAlchemy session object
|
|
"""
|
|
return db_session
|
|
|
|
def get_chroma_collection():
|
|
"""
|
|
Get the ChromaDB collection for Zulip messages.
|
|
|
|
Returns:
|
|
ChromaDB collection object
|
|
"""
|
|
return chroma_collection |