feat: enhance database transactions, add Makefiles, and improve Docker setup
- Refactored database batch operations to ensure transaction atomicity and simplify nested structure. - Added `Makefile` for `knowledge-base` and `llm-gateway` modules to streamline development workflows. - Simplified `Dockerfile` for `llm-gateway` by removing multi-stage builds and optimizing dependencies. - Improved code readability in `collection_manager` and `failover` modules with refined logic. - Minor fixes in `test_server` and Redis health check handling for better diagnostics.
This commit is contained in:
@@ -284,41 +284,40 @@ class DatabaseManager:
|
||||
)
|
||||
|
||||
try:
|
||||
async with self.acquire() as conn:
|
||||
async with self.acquire() as conn, conn.transaction():
|
||||
# Wrap in transaction for all-or-nothing batch semantics
|
||||
async with conn.transaction():
|
||||
for project_id, collection, content, embedding, chunk_type, metadata in embeddings:
|
||||
content_hash = self.compute_content_hash(content)
|
||||
source_path = metadata.get("source_path")
|
||||
start_line = metadata.get("start_line")
|
||||
end_line = metadata.get("end_line")
|
||||
file_type = metadata.get("file_type")
|
||||
for project_id, collection, content, embedding, chunk_type, metadata in embeddings:
|
||||
content_hash = self.compute_content_hash(content)
|
||||
source_path = metadata.get("source_path")
|
||||
start_line = metadata.get("start_line")
|
||||
end_line = metadata.get("end_line")
|
||||
file_type = metadata.get("file_type")
|
||||
|
||||
embedding_id = await conn.fetchval(
|
||||
"""
|
||||
INSERT INTO knowledge_embeddings
|
||||
(project_id, collection, content, embedding, chunk_type,
|
||||
source_path, start_line, end_line, file_type, metadata,
|
||||
content_hash, expires_at)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
|
||||
ON CONFLICT DO NOTHING
|
||||
RETURNING id
|
||||
""",
|
||||
project_id,
|
||||
collection,
|
||||
content,
|
||||
embedding,
|
||||
chunk_type.value,
|
||||
source_path,
|
||||
start_line,
|
||||
end_line,
|
||||
file_type,
|
||||
metadata,
|
||||
content_hash,
|
||||
expires_at,
|
||||
)
|
||||
if embedding_id:
|
||||
ids.append(str(embedding_id))
|
||||
embedding_id = await conn.fetchval(
|
||||
"""
|
||||
INSERT INTO knowledge_embeddings
|
||||
(project_id, collection, content, embedding, chunk_type,
|
||||
source_path, start_line, end_line, file_type, metadata,
|
||||
content_hash, expires_at)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
|
||||
ON CONFLICT DO NOTHING
|
||||
RETURNING id
|
||||
""",
|
||||
project_id,
|
||||
collection,
|
||||
content,
|
||||
embedding,
|
||||
chunk_type.value,
|
||||
source_path,
|
||||
start_line,
|
||||
end_line,
|
||||
file_type,
|
||||
metadata,
|
||||
content_hash,
|
||||
expires_at,
|
||||
)
|
||||
if embedding_id:
|
||||
ids.append(str(embedding_id))
|
||||
|
||||
logger.info(f"Stored {len(ids)} embeddings in batch")
|
||||
return ids
|
||||
@@ -566,59 +565,58 @@ class DatabaseManager:
|
||||
)
|
||||
|
||||
try:
|
||||
async with self.acquire() as conn:
|
||||
async with self.acquire() as conn, conn.transaction():
|
||||
# Use transaction for atomic replace
|
||||
async with conn.transaction():
|
||||
# First, delete existing embeddings for this source
|
||||
delete_result = await conn.execute(
|
||||
# First, delete existing embeddings for this source
|
||||
delete_result = await conn.execute(
|
||||
"""
|
||||
DELETE FROM knowledge_embeddings
|
||||
WHERE project_id = $1 AND source_path = $2 AND collection = $3
|
||||
""",
|
||||
project_id,
|
||||
source_path,
|
||||
collection,
|
||||
)
|
||||
deleted_count = int(delete_result.split()[-1])
|
||||
|
||||
# Then insert new embeddings
|
||||
new_ids = []
|
||||
for content, embedding, chunk_type, metadata in embeddings:
|
||||
content_hash = self.compute_content_hash(content)
|
||||
start_line = metadata.get("start_line")
|
||||
end_line = metadata.get("end_line")
|
||||
file_type = metadata.get("file_type")
|
||||
|
||||
embedding_id = await conn.fetchval(
|
||||
"""
|
||||
DELETE FROM knowledge_embeddings
|
||||
WHERE project_id = $1 AND source_path = $2 AND collection = $3
|
||||
INSERT INTO knowledge_embeddings
|
||||
(project_id, collection, content, embedding, chunk_type,
|
||||
source_path, start_line, end_line, file_type, metadata,
|
||||
content_hash, expires_at)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
|
||||
RETURNING id
|
||||
""",
|
||||
project_id,
|
||||
source_path,
|
||||
collection,
|
||||
content,
|
||||
embedding,
|
||||
chunk_type.value,
|
||||
source_path,
|
||||
start_line,
|
||||
end_line,
|
||||
file_type,
|
||||
metadata,
|
||||
content_hash,
|
||||
expires_at,
|
||||
)
|
||||
deleted_count = int(delete_result.split()[-1])
|
||||
if embedding_id:
|
||||
new_ids.append(str(embedding_id))
|
||||
|
||||
# Then insert new embeddings
|
||||
new_ids = []
|
||||
for content, embedding, chunk_type, metadata in embeddings:
|
||||
content_hash = self.compute_content_hash(content)
|
||||
start_line = metadata.get("start_line")
|
||||
end_line = metadata.get("end_line")
|
||||
file_type = metadata.get("file_type")
|
||||
|
||||
embedding_id = await conn.fetchval(
|
||||
"""
|
||||
INSERT INTO knowledge_embeddings
|
||||
(project_id, collection, content, embedding, chunk_type,
|
||||
source_path, start_line, end_line, file_type, metadata,
|
||||
content_hash, expires_at)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
|
||||
RETURNING id
|
||||
""",
|
||||
project_id,
|
||||
collection,
|
||||
content,
|
||||
embedding,
|
||||
chunk_type.value,
|
||||
source_path,
|
||||
start_line,
|
||||
end_line,
|
||||
file_type,
|
||||
metadata,
|
||||
content_hash,
|
||||
expires_at,
|
||||
)
|
||||
if embedding_id:
|
||||
new_ids.append(str(embedding_id))
|
||||
|
||||
logger.info(
|
||||
f"Replaced source {source_path}: deleted {deleted_count}, "
|
||||
f"inserted {len(new_ids)} embeddings"
|
||||
)
|
||||
return deleted_count, new_ids
|
||||
logger.info(
|
||||
f"Replaced source {source_path}: deleted {deleted_count}, "
|
||||
f"inserted {len(new_ids)} embeddings"
|
||||
)
|
||||
return deleted_count, new_ids
|
||||
|
||||
except asyncpg.PostgresError as e:
|
||||
logger.error(f"Replace source error: {e}")
|
||||
|
||||
Reference in New Issue
Block a user